Compare commits

...
16 Commits
Author SHA1 Message Date
ZyyandCodex 447869860a chore(release): v3.28.0
Co-authored-by: Codex <noreply@openai.com>
2026-10-16 17:15:10 +08:00
ZyyandCodex 815c93a715 test(store): make group capacity fixtures independent of defaults
Set inherited concurrency explicitly so the aggregation test exercises
its intended 3 + 4 capacity without depending on installation defaults.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 17:15:05 +08:00
ZyyandCodex 5843091969 docs: align diagnostic export and troubleshooting guidance
Document server-side redaction of stored error details and timelines,
add the export route, and explain how to inspect diagnostic gaps.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 17:07:59 +08:00
ZyyandCodex 1280c437b4 fix(settings): raise initial account concurrency to five
Set the database default and untouched initial settings to five in the
unreleased 0032 migration, preserving saved concurrency settings.
Update the migration checksum and document the initial value.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 16:59:33 +08:00
ZyyandCodex d965bc9050 fix: harden diagnostics and consolidate unreleased migrations
Redact inline credential headers and token assignments in diagnostic exports.
Remove unnecessary usage query copies and the reverse analytics dependency.

Combine unreleased migrations 0032-0035 without changing published migrations.
Update existing migration fixtures and remove five oversized database stress
tests from the regular suite.

Validate with Store, Provider and diagnostics tests, strict Clippy, Rustfmt,
and frontend lint/build.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 16:55:42 +08:00
bd9b017ed3 feat(usage): unify request records with status filters (#447)
* feat(usage): unify request records with status filters

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

* fix(usage): wrap record filters on narrow screens

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

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-16 16:19:36 +08:00
338cde1505 feat(scheduling): add plugin filtering and remove automatic switchback (#446)
* feat(scheduling): add plugin filtering and remove automatic switchback

Add scheduler v2 filtered delegation across native selection, affinity,
capacity waits, and lease retries. Keep plugin configuration and business
rules owned by plugins without legacy scheduler compatibility.

Remove automatic switchback from the runtime, SDK, and settings UI.
Migrate the existing settings JSON without creating new tables.

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

* fix(dev): align plugin source checks with scheduler v2

Update the workbench submodule to the v2-capable SDK and scheduler,
and document the explicit scheduler manifest version.

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

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-16 15:13:56 +08:00
b25f96559a feat(groups): control image generation by account group (#445)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-16 14:00:14 +08:00
ZyyandCodex a8a96d8919 feat(diagnostics): generate support bundles on the backend
Collect request, retry, error and recovery evidence with bounded reads and
credential redaction. Keep ordinary detail responses small and expose a
shared download action next to the close button.

Preserve export metadata when limits are reached and deduplicate selected
errors without discarding business context based on generic field names.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 12:49:18 +08:00
ZyyandCodex 7b8339ede3 fix(openai): preserve retries after response metadata
Keep response.metadata replayable while forwarding its turn and safety metadata unchanged. Cover delivered metadata before 503 on HTTP/SSE and WebSocket, preserve barriers after output or unknown events, and verify connection-limit recovery.

Co-authored-by: Codex <noreply@openai.com>
2026-10-16 11:16:49 +08:00
f439a30612 feat(usage): expose bounded data maintenance services (#436)
Expose generic detail CRUD, aggregate contributions, and native pricing evaluation through Admin and the plugin SDK. Reuse existing request and statistics storage without schema changes; keep workflow state and recovery with callers.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-16 00:49:51 +08:00
ZyyandCodex 8572873c5f refactor(store): standardize database table names
Rename inconsistent tables and their schema objects, update SQL references,
and preserve analytics trigger behavior through the upgrade.

Co-authored-by: Codex <noreply@openai.com>
2026-10-15 23:54:05 +08:00
57c8ef358f feat(settings): configure request history retention (#435)
Allow 1–31 day retention and keep published analytics readable during cleanup.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-15 22:38:53 +08:00
ZyyandCodex 70ab8ebc8d fix(ui): refine privacy controls and plugin empty state
Co-authored-by: Codex <noreply@openai.com>
2026-10-15 21:34:47 +08:00
ZyyandCodex 14ea5234ce chore(release): v3.27.1
Co-authored-by: Codex <noreply@openai.com>
2026-10-15 18:12:46 +08:00
ZyyandCodex f8480f3bdc fix(privacy): preserve project state in identity presets
Keep repository, commit, branch and worktree data intact while normalizing personal paths. Remove the Desktop identity preset and version shortcuts from privacy settings.

Synchronize automatic Desktop Git context versions from the existing request-scoped client identity in HTTP and WebSocket requests. Reuse official release updates, fixed versions and key overrides, and share context parsing with privacy rules. Preserve unknown or incomplete context and leave installation identity to its existing account owner.

Verified version field semantics against the official Desktop 26.1007.71326 bundle and Codex app-server initialization source. Validated provider tests, temporary HTTP/WS identity tests, architecture boundaries, strict Clippy, frontend lint/build and the preset UI.

Co-authored-by: Codex <noreply@openai.com>
2026-10-15 16:45:04 +08:00
223 changed files with 8626 additions and 2359 deletions
+1
View File
@@ -1746,6 +1746,7 @@ dependencies = [
"http",
"http-body",
"http-body-util",
"regex",
"secrecy",
"serde",
"serde_json",
@@ -107,7 +107,7 @@ async fn execute(
let settings = gateway_admin::initialize_settings(
store.admin_ports().settings(),
core.snapshot_control(),
providers.admin,
providers.admin.clone(),
std::sync::Arc::new(gateway_host::pricing::ModelsDevPricing),
);
let mut services = gateway_admin::public_service::Registry::new(service_middleware(
@@ -115,6 +115,12 @@ async fn execute(
&plugin_runtime,
));
or_shutdown!(plugin_runtime, services.register_settings(&settings));
let usage_data = gateway_admin::initialize_usage_data(
store.admin_ports().usage_data(),
store.admin_ports().settings(),
providers.admin,
);
or_shutdown!(plugin_runtime, services.register_usage_data(&usage_data));
let services = std::sync::Arc::new(services);
or_shutdown!(plugin_runtime, plugin_runtime.bind_services(&services));
let commands = or_shutdown!(plugin_runtime, plugin_runtime.prepare_command_line().await);
+23
View File
@@ -39,6 +39,7 @@ pub use use_case::{
proxies::ProxiesService,
settings::SettingsService,
system::SystemService,
usage_data::UsageDataService,
};
use model::AdminError;
@@ -217,6 +218,7 @@ pub struct AdminServices {
client_distribution: Arc<dyn ClientDistributionService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
usage_data: Arc<dyn UsageDataService>,
system: Arc<dyn SystemService>,
credentials: Arc<CredentialsService>,
// Runtime 只持有 Weak,这个强引用保障账号端口与 Admin 服务同寿命
@@ -285,6 +287,11 @@ impl AdminServices {
self.client_distribution.as_ref()
}
#[must_use]
pub fn usage_data(&self) -> &dyn UsageDataService {
self.usage_data.as_ref()
}
#[must_use]
pub fn observability(&self) -> &dyn ObservabilityService {
self.observability.as_ref()
@@ -467,8 +474,10 @@ async fn initialize_inner(
registry.clone(),
pricing_source,
);
let usage_data = initialize_usage_data(store.usage_data(), store.settings(), registry.clone());
let mut public_services = public_service::Registry::new(service_middleware);
public_services.register_settings(&settings)?;
public_services.register_usage_data(&usage_data)?;
let services = AdminServices {
timezone,
public_services: Arc::new(public_services),
@@ -512,8 +521,10 @@ async fn initialize_inner(
store.settings(),
registry.clone(),
timezone,
system.clone(),
)),
settings,
usage_data,
system,
credentials,
_plugin_accounts: plugin_accounts,
@@ -644,3 +655,15 @@ pub fn initialize_settings(
pricing_source,
))
}
/// HTTP、插件及 CLI 共用的用量领域用例,宿主提供已验证的调用主体
#[must_use]
pub fn initialize_usage_data(
store: Arc<dyn ports::usage_data::UsageDataStore>,
settings: Arc<dyn ports::store::SettingsStore>,
providers: ProviderAdminRegistry,
) -> Arc<dyn UsageDataService> {
Arc::new(use_case::usage_data::DefaultUsageDataService::new(
store, settings, providers,
))
}
@@ -100,6 +100,7 @@ pub struct AccountGroupListQuery {
/// 完整账号分组摘要
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupRecord {
pub allow_image_generation: bool,
pub fast_mode: FastMode,
pub id: AccountGroupId,
pub name: String,
@@ -139,6 +140,7 @@ pub struct AccountGroupOptionsPage {
/// 创建账号分组
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateAccountGroup {
pub allow_image_generation: bool,
pub fast_mode: FastMode,
pub name: String,
pub description: Option<String>,
@@ -148,6 +150,7 @@ pub struct CreateAccountGroup {
/// 已生成稳定标识、可直接提交存储的创建命令
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewAccountGroup {
pub allow_image_generation: bool,
pub fast_mode: FastMode,
pub id: AccountGroupId,
pub name: String,
@@ -158,6 +161,7 @@ pub struct NewAccountGroup {
/// 更新账号分组的描述字段
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateAccountGroup {
pub allow_image_generation: Option<bool>,
pub fast_mode: Option<FastMode>,
pub id: AccountGroupId,
pub name: String,
@@ -0,0 +1,44 @@
//! 故障诊断导出的目标、已采集证据与缺口
use chrono::{DateTime, Utc};
use super::{
observability::{OpsError, UsageDetail},
system::SystemVersion,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DiagnosticTarget {
Request(String),
Operation(String),
}
impl DiagnosticTarget {
#[must_use]
pub fn id(&self) -> &str {
match self {
Self::Request(id) | Self::Operation(id) => id,
}
}
}
pub struct DiagnosticGap {
pub section: String,
pub reason: &'static str,
}
pub struct DiagnosticRequest {
pub request_id: String,
pub detail: Option<UsageDetail>,
pub errors: Vec<OpsError>,
}
/// 原始受控详情只能经导出投影脱敏后交付,不参与普通日志或 Debug
pub struct DiagnosticBundle {
pub target: DiagnosticTarget,
pub collected_at: DateTime<Utc>,
pub environment: Option<SystemVersion>,
pub error: Option<OpsError>,
pub requests: Vec<DiagnosticRequest>,
pub gaps: Vec<DiagnosticGap>,
}
@@ -10,6 +10,7 @@ pub mod auth;
pub mod backup;
pub mod client_distribution;
pub mod client_keys;
pub mod diagnostics;
pub mod import_tasks;
pub mod key_usage;
pub mod observability;
@@ -24,6 +25,7 @@ pub mod quota_forecast_sampling;
pub mod retention;
pub mod settings;
pub mod system;
pub mod usage_data;
/// 管理用例对外返回的稳定错误分类
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -172,9 +172,19 @@ pub struct UsageFilter {
pub search: Option<String>,
}
/// 请求列表的结果筛选,不改变汇总和计量口径
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum UsageRecordStatus {
#[default]
Success,
Error,
All,
}
/// 用量记录分页查询
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageQuery {
pub status: UsageRecordStatus,
pub range: TimeRange,
pub filter: UsageFilter,
pub current_page: u32,
@@ -653,6 +663,7 @@ pub struct DashboardObservation {
/// 使用记录表格的窄读模型
#[derive(Debug, Clone, PartialEq)]
pub struct UsageListRecord {
pub outcome: RequestOutcome,
pub client_api_key_name: Option<String>,
pub id: String,
pub endpoint: String,
@@ -845,7 +856,7 @@ pub struct UsageAttempt {
pub occurred_at: DateTime<Utc>,
}
/// 一条请求及其全部尝试
/// 一条请求及其保留的尝试与诊断轨迹
#[derive(Debug, Clone, PartialEq)]
pub struct UsageDetail {
pub trace: Option<serde_json::Value>,
@@ -24,7 +24,7 @@ impl RetentionPolicy {
audit_days: u32,
) -> Result<(), &'static str> {
for (valid, field) in [
(usage_days >= 31, "usage_retention_days"),
((1..=31).contains(&usage_days), "usage_retention_days"),
(ops_days > 0, "ops_event_retention_days"),
(audit_days > 0, "audit_retention_days"),
] {
@@ -38,7 +38,7 @@ impl RetentionPolicy {
pub fn try_new(usage_days: u32, ops_days: u32, audit_days: u32) -> Result<Self, AdminError> {
if Self::validate_values(usage_days, ops_days, audit_days).is_err() {
return Err(AdminError::invalid(
"请求保留期至少 31 天,事件与审计保留期必须为正数",
"请求保留期必须为 1~31 天,事件与审计保留期必须为正数",
));
}
Ok(Self {
@@ -0,0 +1,26 @@
//! 由宿主传入的可信写入主体,不属于插件请求载荷
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsageActor {
Admin {
subject: String,
},
Plugin {
publisher: String,
name: String,
instance_id: String,
},
System,
}
impl super::UsageRecordChange {
#[must_use]
pub fn id(&self) -> &str {
match self {
Self::Put { id, .. } | Self::Update { id, .. } | Self::Delete { id } => id,
}
}
}
@@ -0,0 +1,87 @@
//! 按归属时间和维度维护统计贡献,业务来源与调整历史由调用方保存
use super::{UsageDelta, UsageRecordFilter};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageMetric {
RequestCount,
SuccessCount,
FailureCount,
CancelledCount,
IncompleteCount,
CompletedRequestCount,
InputTokens,
OutputTokens,
CachedInputTokens,
CacheWriteInputTokens,
ReasoningOutputTokens,
ImageInputTokens,
ImageOutputTokens,
TotalTokens,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageDimensions {
pub provider: Option<String>,
pub model: Option<String>,
pub account_ref: Option<String>,
pub client_key_ref: Option<String>,
pub operation: Option<String>,
pub transport: Option<String>,
}
/// 整个键决定覆盖目标;调用方可将来源标识编码进 key,重试相同键不会累计
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageAggregateKey {
pub bucket_start_ms: i64,
pub dimensions: UsageDimensions,
pub key: String,
}
/// 值按 bucket_start_ms 归属,不自动摊分;负值用于相对原生统计的修正
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageAggregate {
pub identity: UsageAggregateKey,
pub counts: BTreeMap<UsageMetric, UsageDelta>,
pub costs: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsageAggregateChange {
Put { aggregate: UsageAggregate },
Delete { identity: UsageAggregateKey },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WriteUsageAggregates {
pub changes: Vec<UsageAggregateChange>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageAggregateWriteResult {
pub changed: Vec<UsageAggregateKey>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ListUsageAggregates {
pub filter: UsageRecordFilter,
pub cursor: Option<String>,
pub limit: Option<u16>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageAggregatePage {
pub items: Vec<UsageAggregate>,
pub next_cursor: Option<String>,
}
@@ -0,0 +1,44 @@
//! 用量基础操作的错误分类,不保存失败或重试状态
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageErrorKind {
InvalidInput,
NotFound,
Conflict,
CapacityExceeded,
Unavailable,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageDataError {
pub kind: UsageErrorKind,
pub message: String,
}
impl UsageDataError {
#[must_use]
pub fn new(kind: UsageErrorKind, message: impl Into<String>) -> Self {
Self {
kind,
message: message.into(),
}
}
#[must_use]
pub fn invalid(message: impl Into<String>) -> Self {
Self::new(UsageErrorKind::InvalidInput, message)
}
#[must_use]
pub fn unavailable(message: impl Into<String>) -> Self {
Self::new(UsageErrorKind::Unavailable, message)
}
}
impl std::fmt::Display for UsageDataError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for UsageDataError {}
pub type UsageDataResult<T> = Result<T, UsageDataError>;
@@ -0,0 +1,22 @@
//! 用量明细、聚合维护与纯计价的公开合同
mod value;
pub use value::*;
mod record;
pub use record::*;
mod query;
pub use query::*;
mod pricing;
pub use pricing::*;
mod error;
pub use error::*;
pub const USAGE_DATA_MAX_TARGETS: usize = 500;
pub const USAGE_DATA_MAX_REQUEST_BYTES: usize = 2 * 1024 * 1024;
pub const USAGE_DATA_MAX_PAGE_SIZE: u16 = 200;
mod actor;
pub use actor::*;
mod aggregate;
pub use aggregate::*;
@@ -0,0 +1,62 @@
//! 纯 Token 计价合同,费用是否采用与额外费用合并由调用方决定
use super::{UsageCost, UsageTokens};
use gateway_core::metering::ModelPriceOverride;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsagePriceSelection {
Current,
Override { price: Box<ModelPriceOverride> },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsagePricingInput {
pub id: String,
pub provider: String,
pub model: String,
pub tokens: UsageTokens,
pub service_tier: Option<String>,
pub price: UsagePriceSelection,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EvaluateUsagePricing {
pub items: Vec<UsagePricingInput>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsagePricingBreakdown {
pub input_amount: String,
pub output_amount: String,
pub cache_read_amount: String,
pub cache_write_amount: String,
pub standard_amount: String,
pub long_context_billing_applied: bool,
pub service_tier: Option<String>,
pub custom_multiplier_bps: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsagePricingResult {
Estimated {
id: String,
cost: UsageCost,
breakdown: UsagePricingBreakdown,
},
NotEstimable {
id: String,
reason: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsagePricingEvaluation {
pub items: Vec<UsagePricingResult>,
}
@@ -0,0 +1,36 @@
//! 请求明细的有界筛选与稳定顺序分页
use super::UsageRecord;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GetUsageRecord {
pub id: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageRecordFilter {
pub start_ms: i64,
pub end_ms: i64,
pub provider: Option<String>,
pub model: Option<String>,
pub account_ref: Option<String>,
pub client_key_ref: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ListUsageRecords {
pub filter: UsageRecordFilter,
pub cursor: Option<String>,
pub limit: Option<u16>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageRecordPage {
pub items: Vec<UsageRecord>,
pub next_cursor: Option<String>,
}
@@ -0,0 +1,106 @@
//! 既有请求明细的用量读写合同,不承载插件任务或来源状态
use super::UsageCount;
use serde::{Deserialize, Serialize};
/// Token 使用宿主口径:缓存包含在输入中,推理包含在输出中,未知值为 null
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageTokens {
pub input_tokens: Option<UsageCount>,
pub output_tokens: Option<UsageCount>,
pub cached_input_tokens: Option<UsageCount>,
pub cache_write_input_tokens: Option<UsageCount>,
pub reasoning_output_tokens: Option<UsageCount>,
pub image_input_tokens: Option<UsageCount>,
pub image_output_tokens: Option<UsageCount>,
pub total_tokens: Option<UsageCount>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageCostSource {
ProviderReported,
Calculated,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsageCost {
#[default]
Unavailable,
Known {
source: UsageCostSource,
amount: String,
currency: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageOutcome {
Running,
Succeeded,
Failed,
Cancelled,
Incomplete,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageRecordData {
pub client_key_ref: String,
pub provider: Option<String>,
pub account_ref: Option<String>,
pub requested_model: Option<String>,
pub model: Option<String>,
pub protocol: String,
pub operation: String,
pub endpoint: String,
pub client_transport: String,
pub service_tier: Option<String>,
pub started_at_ms: i64,
pub completed_at_ms: Option<i64>,
pub outcome: UsageOutcome,
pub tokens: UsageTokens,
pub cost: UsageCost,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageRecord {
pub id: String,
pub data: UsageRecordData,
pub billing_snapshot: Option<serde_json::Value>,
pub provider_observation: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum UsageRecordChange {
Put {
id: String,
data: Box<UsageRecordData>,
},
/// tokens 替换全部 Token 字段;cost 的 unavailable 明确清空费用
Update {
id: String,
tokens: Option<UsageTokens>,
cost: Option<UsageCost>,
},
Delete {
id: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WriteUsageRecords {
pub changes: Vec<UsageRecordChange>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UsageWriteResult {
pub changed_ids: Vec<String>,
}
@@ -0,0 +1,63 @@
//! 用量计数使用十进制字符串,避免插件前端丢失整数精度
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error};
use std::fmt;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UsageCount(pub u64);
impl UsageCount {
pub const ZERO: Self = Self(0);
#[must_use]
pub const fn get(self) -> u64 {
self.0
}
}
impl From<u64> for UsageCount {
fn from(value: u64) -> Self {
Self(value)
}
}
impl fmt::Display for UsageCount {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
impl Serialize for UsageCount {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(&self.0.to_string())
}
}
impl<'de> Deserialize<'de> for UsageCount {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let value = String::deserialize(deserializer)?;
if value.is_empty() || !value.bytes().all(|byte| byte.is_ascii_digit()) {
return Err(D::Error::custom("count must be an unsigned decimal string"));
}
value.parse().map(Self).map_err(D::Error::custom)
}
}
/// 有符号统计贡献受 PostgreSQL bigint 范围限制,JSON 使用十进制字符串
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UsageDelta(pub i64);
impl UsageDelta {
#[must_use]
pub const fn get(self) -> i64 {
self.0
}
}
impl Serialize for UsageDelta {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(&self.0.to_string())
}
}
impl<'de> Deserialize<'de> for UsageDelta {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let value = String::deserialize(deserializer)?;
let digits = value.strip_prefix('-').unwrap_or(&value);
if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) {
return Err(D::Error::custom("delta must be a signed decimal string"));
}
value.parse().map(Self).map_err(D::Error::custom)
}
}
@@ -14,3 +14,4 @@ pub mod proxy;
pub mod retention;
pub mod store;
pub mod system;
pub mod usage_data;
@@ -26,6 +26,14 @@ use crate::model::{
quota_forecast_sampling::QuotaForecastObservation,
};
/// 按原生 Token 口径计算的报价输入,额外费用由调用方处理
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderUsagePricingInput {
pub upstream_model_id: String,
pub service_tier: Option<String>,
pub usage: gateway_core::metering::Usage,
}
/// Provider 管理失败的稳定分类
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderAdminErrorKind {
@@ -161,6 +169,14 @@ pub trait ProviderAdmin: Send + Sync {
Default::default()
}
fn evaluate_usage_pricing(
&self,
_input: &ProviderUsagePricingInput,
_pricing: Option<&gateway_core::metering::ModelPriceOverride>,
) -> Result<Option<gateway_core::metering::CalculatedCostBreakdown>, ProviderAdminError> {
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
}
fn provider_kind(&self) -> &ProviderKind;
/// 只读取原生 Provider 能力及账号类型,不执行网络或数据库查询
@@ -485,6 +485,13 @@ pub trait ObservabilityStore: Send + Sync {
) -> AdminStoreResult<DiagnosticsObservation>;
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage>;
/// 按精确目标读取受控错误,不受列表时间窗口与当前页面影响
async fn diagnostic_errors(
&self,
target: &crate::model::diagnostics::DiagnosticTarget,
limit: u16,
) -> AdminStoreResult<Vec<crate::model::observability::OpsError>>;
}
/// Runtime settings 与管理员 API Key 写入
@@ -565,9 +572,15 @@ pub struct AdminStorePorts {
plugins: Arc<dyn super::plugins::PluginStore>,
plugin_state: Arc<dyn super::plugins::PluginStateStore>,
plugin_resources: Arc<dyn super::plugin_resources::PluginResourceStore>,
usage_data: Arc<dyn super::usage_data::UsageDataStore>,
}
impl AdminStorePorts {
#[must_use]
pub fn usage_data(&self) -> Arc<dyn super::usage_data::UsageDataStore> {
self.usage_data.clone()
}
#[must_use]
pub fn plugin_resources(&self) -> Arc<dyn super::plugin_resources::PluginResourceStore> {
self.plugin_resources.clone()
@@ -598,6 +611,7 @@ impl AdminStorePorts {
plugins: Arc<dyn super::plugins::PluginStore>,
plugin_state: Arc<dyn super::plugins::PluginStateStore>,
plugin_resources: Arc<dyn super::plugin_resources::PluginResourceStore>,
usage_data: Arc<dyn super::usage_data::UsageDataStore>,
) -> Self {
Self {
accounts,
@@ -609,6 +623,7 @@ impl AdminStorePorts {
plugins,
plugin_state,
plugin_resources,
usage_data,
}
}
@@ -0,0 +1,27 @@
//! 原生请求用量的有界读取和事务写入端口
use crate::model::{auth::AdminAuditEvent, usage_data::*};
use async_trait::async_trait;
pub struct StoredUsageWrite {
pub changes: Vec<(UsageRecordChange, AdminAuditEvent)>,
}
pub struct StoredAggregateWrite {
pub changes: Vec<(UsageAggregateChange, AdminAuditEvent)>,
}
#[async_trait]
pub trait UsageDataStore: Send + Sync {
async fn list_aggregates(
&self,
request: ListUsageAggregates,
) -> UsageDataResult<UsageAggregatePage>;
async fn write_aggregates(
&self,
command: StoredAggregateWrite,
) -> UsageDataResult<UsageAggregateWriteResult>;
async fn get_record(&self, request: GetUsageRecord) -> UsageDataResult<UsageRecord>;
async fn list_records(&self, request: ListUsageRecords) -> UsageDataResult<UsageRecordPage>;
async fn write_records(&self, command: StoredUsageWrite) -> UsageDataResult<UsageWriteResult>;
}
@@ -0,0 +1,21 @@
//! 公开服务写入主体由受信宿主边界设置,插件输入不能覆盖
use crate::model::usage_data::UsageActor;
use std::future::Future;
tokio::task_local! {
static ACTOR: UsageActor;
}
/// 与 middleware Origin 独立传递,嵌套服务不会丢失真实调用主体
pub async fn scope_actor<F: Future>(actor: UsageActor, future: F) -> F::Output {
ACTOR.scope(actor, future).await
}
pub(super) fn current_actor() -> Result<UsageActor, crate::model::usage_data::UsageDataError> {
ACTOR.try_with(Clone::clone).map_err(|_| {
crate::model::usage_data::UsageDataError::invalid(
"usage mutation requires a trusted caller identity",
)
})
}
@@ -1,7 +1,10 @@
//! 主动服务调用的统一组合入口;普通业务方法由所属 HTTP 或执行边界包裹
mod actor;
mod registry;
mod settings;
mod usage_data;
pub use actor::scope_actor;
pub use registry::Registry;
use std::{future::Future, sync::Arc};
@@ -54,6 +54,57 @@ impl Registry {
Ok(())
}
pub fn register_usage_data(
&mut self,
service: &Arc<dyn crate::UsageDataService>,
) -> Result<(), AdminError> {
super::usage_data::register(self, service)
}
pub(crate) fn register_usage<I, O>(
&mut self,
name: &'static str,
handler: impl Fn(I) -> BoxFuture<'static, crate::model::usage_data::UsageDataResult<O>>
+ Send
+ Sync
+ 'static,
) -> Result<(), AdminError>
where
I: DeserializeOwned + Send + 'static,
O: Serialize + Send + 'static,
{
let std::collections::btree_map::Entry::Vacant(entry) = self.operations.entry(name) else {
return Err(AdminError::internal(format!("公开服务重复登记:{name}")));
};
let handler = Arc::new(handler);
entry.insert(Arc::new(move |input| {
let handler = handler.clone();
Box::pin(async move {
// 在反序列化之前限制编码后的完整请求,避免单条记录绕过批量行数约束
let bytes =
serde_json::to_vec(&input).map_err(|_| Error::invalid("服务输入编码无效"))?;
if bytes.len() > crate::model::usage_data::USAGE_DATA_MAX_REQUEST_BYTES {
return Err(super::usage_data::encode_error(
crate::model::usage_data::UsageDataError::new(
crate::model::usage_data::UsageErrorKind::CapacityExceeded,
"usage input exceeds 2 MiB",
),
));
}
let input = serde_json::from_value(input).map_err(|_| {
super::usage_data::encode_error(
crate::model::usage_data::UsageDataError::invalid("服务输入类型无效"),
)
})?;
let output = handler(input)
.await
.map_err(super::usage_data::encode_error)?;
serde_json::to_value(output).map_err(|_| Error::unavailable("服务结果编码失败"))
})
}));
Ok(())
}
pub async fn call(&self, operation: &str, input: Value) -> Result<Value, Error> {
let (operation, handler) =
self.operations
@@ -0,0 +1,63 @@
//! 原生用量能力的公开服务登记
use crate::{
UsageDataService,
model::{AdminError, usage_data::*},
};
use std::sync::Arc;
pub(super) fn register(
registry: &mut super::Registry,
service: &Arc<dyn UsageDataService>,
) -> Result<(), AdminError> {
macro_rules! read {
($name:literal, $method:ident, $input:ty) => {{
let service = service.clone();
registry.register_usage($name, move |request: $input| {
let service = service.clone();
Box::pin(async move { service.$method(request).await })
})?;
}};
}
read!("usage.records.get", get_record, GetUsageRecord);
read!("usage.records.list", list_records, ListUsageRecords);
read!(
"usage.aggregates.list",
list_aggregates,
ListUsageAggregates
);
read!("pricing.evaluate", evaluate_pricing, EvaluateUsagePricing);
let aggregates = service.clone();
registry.register_usage(
"usage.aggregates.write",
move |request: WriteUsageAggregates| {
let service = aggregates.clone();
Box::pin(async move {
let actor = super::actor::current_actor()?;
service.write_aggregates(&actor, request).await
})
},
)?;
let service = service.clone();
registry.register_usage("usage.records.write", move |request: WriteUsageRecords| {
let service = service.clone();
Box::pin(async move {
let actor = super::actor::current_actor()?;
service.write_records(&actor, request).await
})
})?;
Ok(())
}
pub(super) fn encode_error(
error: UsageDataError,
) -> gateway_core::engine::middleware::service::Error {
gateway_core::engine::middleware::service::Error {
kind: serde_json::to_value(error.kind)
.ok()
.and_then(|v| v.as_str().map(str::to_owned))
.unwrap_or_else(|| "unavailable".into()),
message: error.message.clone(),
details: serde_json::to_value(error).ok(),
}
}
@@ -160,6 +160,7 @@ impl AccountGroupService for DefaultAccountGroupService {
self.store
.create_account_group(
NewAccountGroup {
allow_image_generation: command.allow_image_generation,
id,
name: command.name,
description: command.description,
@@ -225,6 +225,7 @@ impl KeyUsageService for DefaultKeyUsageService {
KeyUsageRecordKind::Success => KeyUsageRecords::Success(
self.observations
.list_usage_records(UsageQuery {
status: Default::default(),
range: query.usage.range,
filter: usage_filter(&id, query.usage.model),
current_page: query.current_page,
@@ -18,6 +18,7 @@ pub mod plugins;
pub mod proxies;
pub mod settings;
pub mod system;
pub mod usage_data;
mod credential_mutation;
mod error;
@@ -0,0 +1,227 @@
//! 按故障目标收集诊断证据,补充读取失败不丢弃已取得的记录
use std::{
collections::{HashSet, VecDeque},
time::Duration,
};
use tokio::time::{Instant, timeout_at};
use super::{DefaultObservabilityService, map_store_error};
use crate::model::{
AdminError, AdminErrorKind,
diagnostics::{DiagnosticBundle, DiagnosticGap, DiagnosticRequest, DiagnosticTarget},
observability::OpsError,
};
const MAX_REQUESTS: usize = 8;
const MAX_ERRORS: u16 = 128;
impl DefaultObservabilityService {
pub(super) async fn collect_diagnostics(
&self,
target: DiagnosticTarget,
) -> Result<DiagnosticBundle, AdminError> {
if target.id().trim().is_empty() || target.id().len() > 256 {
return Err(AdminError::invalid("诊断目标 ID 不合法"));
}
let deadline = Instant::now() + Duration::from_secs(20);
let (collected, environment) = tokio::join!(
self.collect_records(target, deadline),
timeout_at(deadline, self.system.version()),
);
let mut bundle = collected?;
match environment {
Ok(Ok(version)) => bundle.environment = Some(version),
_ => gap(&mut bundle, "environment", "collection_failed"),
}
Ok(bundle)
}
async fn collect_records(
&self,
target: DiagnosticTarget,
deadline: Instant,
) -> Result<DiagnosticBundle, AdminError> {
let mut bundle = DiagnosticBundle {
target,
collected_at: chrono::Utc::now(),
environment: None,
error: None,
requests: Vec::new(),
gaps: Vec::new(),
};
let mut remaining_errors = MAX_ERRORS;
let first_id = match &bundle.target {
DiagnosticTarget::Request(id) => Some(id.clone()),
DiagnosticTarget::Operation(_) => {
let mut errors = self.export_errors(&bundle.target, 1, deadline).await?;
let error = errors
.pop()
.ok_or_else(|| AdminError::not_found("错误记录不存在或已过期"))?;
let id = error.request_id.clone();
bundle.error = Some(error);
remaining_errors -= 1;
id
}
};
let mut pending = VecDeque::from_iter(first_id);
let mut seen = HashSet::new();
while let Some(id) = pending.pop_front() {
if !seen.insert(id.clone()) {
continue;
}
if bundle.requests.len() >= MAX_REQUESTS {
gap(&mut bundle, "requests", "related_request_limit");
break;
}
let section = format!("requests.{id}");
let target = DiagnosticTarget::Request(id.clone());
let selected = bundle
.error
.as_ref()
.filter(|error| error.request_id.as_deref() == Some(id.as_str()));
// 事件导出已经持有选中记录,额外读取一个槽位用于去重后判断是否超限
let read_limit = remaining_errors + 1 + u16::from(selected.is_some());
let (detail, errors) = tokio::join!(
timeout_at(deadline, self.store.usage_record_detail(&id)),
self.export_errors(&target, read_limit, deadline),
);
let (detail, detail_error) = match detail {
Ok(Ok(detail)) => (Some(detail), None),
Ok(Err(error)) => (None, Some(map_store_error(error, "diagnostic request"))),
Err(_) => (
None,
Some(AdminError::new(AdminErrorKind::Unavailable, "诊断采集超时")),
),
};
let mut errors = match errors {
Ok(mut errors) => {
if let Some(selected) = selected {
errors.retain(|error| {
error.source != selected.source || error.event_id != selected.event_id
});
}
errors
}
Err(error) => {
if detail.is_none() && bundle.requests.is_empty() && bundle.error.is_none() {
return Err(error);
}
gap(
&mut bundle,
&format!("{section}.errors"),
"collection_failed",
);
Vec::new()
}
};
if detail.is_none()
&& errors.is_empty()
&& bundle.requests.is_empty()
&& bundle.error.is_none()
{
return Err(
detail_error.unwrap_or_else(|| AdminError::not_found("请求记录不存在或已过期"))
);
}
if errors.len() > usize::from(remaining_errors) {
gap(
&mut bundle,
&format!("{section}.errors"),
"error_record_limit",
);
errors.truncate(usize::from(remaining_errors));
}
remaining_errors =
remaining_errors.saturating_sub(u16::try_from(errors.len()).unwrap_or(MAX_ERRORS));
if let Some(record) = &detail {
for related in &record.related_requests {
if let Some(id) = related.get("requestId").and_then(serde_json::Value::as_str) {
pending.push_back(id.to_owned());
}
}
if record.request.outcome == crate::model::observability::RequestOutcome::Running {
gap(&mut bundle, &section, "request_not_finalized");
}
match &record.trace {
None => gap(
&mut bundle,
&format!("{section}.trace"),
"not_recorded_or_retained",
),
Some(trace) => {
if trace["droppedEvents"].as_u64().unwrap_or(0) > 0 {
gap(&mut bundle, &format!("{section}.trace"), "events_evicted");
}
let events = trace["events"].as_array();
if !events.is_some_and(|events| {
events.iter().any(|event| {
event["stage"] == "request.routed"
&& event["data"]["accountPolicy"].is_object()
})
}) {
gap(&mut bundle, &section, "incident_policy_not_recorded");
}
for event in events.into_iter().flatten() {
if event["stage"] == "capture.gap" || event["data"]["truncated"] == true
{
gap(
&mut bundle,
&format!("{section}.trace"),
"capture_gap_or_truncation",
);
}
if event["stage"] == "retry.decided" && event["data"]["phase"].is_null()
{
gap(
&mut bundle,
&format!("{section}.trace"),
"legacy_retry_facts_not_recorded",
);
}
}
}
}
gap(
&mut bundle,
&format!("{section}.attempts"),
"best_effort_persistence",
);
} else {
gap(&mut bundle, &section, "request_detail_unavailable");
}
bundle.requests.push(DiagnosticRequest {
request_id: id,
detail,
errors,
});
}
Ok(bundle)
}
async fn export_errors(
&self,
target: &DiagnosticTarget,
limit: u16,
deadline: Instant,
) -> Result<Vec<OpsError>, AdminError> {
timeout_at(deadline, self.store.diagnostic_errors(target, limit))
.await
.map_err(|_| AdminError::new(AdminErrorKind::Unavailable, "诊断采集超时"))?
.map_err(|error| map_store_error(error, "diagnostic errors"))
}
}
fn gap(bundle: &mut DiagnosticBundle, section: &str, reason: &'static str) {
if !bundle
.gaps
.iter()
.any(|gap| gap.section == section && gap.reason == reason)
{
bundle.gaps.push(DiagnosticGap {
section: section.to_owned(),
reason,
});
}
}
@@ -38,6 +38,7 @@ const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10;
const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3;
const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 95.0;
mod export;
mod store_cache;
pub(crate) use store_cache::CachedObservabilityStore;
@@ -70,6 +71,10 @@ pub trait ObservabilityService: Send + Sync {
dimension: DiagnosticDimension,
) -> Result<DiagnosticsResult, AdminError>;
async fn ops_errors(&self, query: OpsErrorQuery) -> Result<OpsErrorPage, AdminError>;
async fn export_diagnostics(
&self,
target: crate::model::diagnostics::DiagnosticTarget,
) -> Result<crate::model::diagnostics::DiagnosticBundle, AdminError>;
}
pub(crate) struct DefaultObservabilityService {
@@ -78,6 +83,7 @@ pub(crate) struct DefaultObservabilityService {
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
system: Arc<dyn super::system::SystemService>,
}
impl DefaultObservabilityService {
@@ -88,6 +94,7 @@ impl DefaultObservabilityService {
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
timezone: gateway_core::time::DeploymentTimeZone,
system: Arc<dyn super::system::SystemService>,
) -> Self {
Self {
timezone,
@@ -95,6 +102,7 @@ impl DefaultObservabilityService {
accounts,
settings,
providers,
system,
}
}
@@ -192,6 +200,13 @@ impl DefaultObservabilityService {
#[async_trait]
impl ObservabilityService for DefaultObservabilityService {
async fn export_diagnostics(
&self,
target: crate::model::diagnostics::DiagnosticTarget,
) -> Result<crate::model::diagnostics::DiagnosticBundle, AdminError> {
self.collect_diagnostics(target).await
}
async fn dashboard_summary(
&self,
range: TimeRange,
@@ -1,4 +1,4 @@
//! 有界统计缓存:身份和完整筛选参与键,删除版本阻止旧数据复活
//! 有界统计缓存:身份和完整筛选参与键,统计发布代次使旧快照失效
use crate::model::observability::*;
use crate::ports::store::{
@@ -121,7 +121,7 @@ impl<K: Eq + Hash, T: AnalyticsValue> Cache<K, T> {
};
// 同键刷新只执行一次;失败不会缓存成成功,也不延长旧结果寿命
let mut cached = slot.lock().await;
// 等锁期间代次可能推进;取得同键锁后再核对,失败时也不能回用已删除事实
// 等锁期间代次可能推进;取得同键锁后再核对,失败时也不能回用已被新统计替代的快照
let generation = store.analytics_cache_generation().await?;
if let Some(entry) = cached.as_mut()
&& entry.value.metadata().generation == generation
@@ -289,4 +289,12 @@ impl ObservabilityStore for CachedObservabilityStore {
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage> {
self.inner.list_ops_errors(query).await
}
async fn diagnostic_errors(
&self,
target: &crate::model::diagnostics::DiagnosticTarget,
limit: u16,
) -> AdminStoreResult<Vec<crate::model::observability::OpsError>> {
self.inner.diagnostic_errors(target, limit).await
}
}
@@ -52,6 +52,7 @@ impl PluginResourceAccess for DefaultPluginResourceAccess {
owner,
resource_key,
NewAccountGroup {
allow_image_generation: command.allow_image_generation,
id,
name: command.name,
description: command.description,
@@ -0,0 +1,184 @@
//! 公开原生用量读写与纯计价,插件拥有跨调用流程和恢复策略
mod pricing;
mod validation;
use crate::{
model::{
auth::{AdminAuditEvent, AuditActorKind},
usage_data::*,
},
ports::{provider::ProviderAdminRegistry, store::SettingsStore, usage_data::*},
};
use async_trait::async_trait;
use std::sync::Arc;
#[async_trait]
pub trait UsageDataService: Send + Sync {
async fn list_aggregates(
&self,
request: ListUsageAggregates,
) -> UsageDataResult<UsageAggregatePage>;
async fn write_aggregates(
&self,
actor: &UsageActor,
request: WriteUsageAggregates,
) -> UsageDataResult<UsageAggregateWriteResult>;
async fn get_record(&self, request: GetUsageRecord) -> UsageDataResult<UsageRecord>;
async fn list_records(&self, request: ListUsageRecords) -> UsageDataResult<UsageRecordPage>;
async fn write_records(
&self,
actor: &UsageActor,
request: WriteUsageRecords,
) -> UsageDataResult<UsageWriteResult>;
async fn evaluate_pricing(
&self,
request: EvaluateUsagePricing,
) -> UsageDataResult<UsagePricingEvaluation>;
}
pub struct DefaultUsageDataService {
store: Arc<dyn UsageDataStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
}
impl DefaultUsageDataService {
pub fn new(
store: Arc<dyn UsageDataStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
) -> Self {
Self {
store,
settings,
providers,
}
}
}
#[async_trait]
impl UsageDataService for DefaultUsageDataService {
async fn list_aggregates(
&self,
request: ListUsageAggregates,
) -> UsageDataResult<UsageAggregatePage> {
validation::query(&request.filter, request.cursor.as_deref(), request.limit)?;
self.store.list_aggregates(request).await
}
async fn write_aggregates(
&self,
actor: &UsageActor,
request: WriteUsageAggregates,
) -> UsageDataResult<UsageAggregateWriteResult> {
validation::bounded(&request, request.changes.len())?;
let mut changes = Vec::with_capacity(request.changes.len());
for change in request.changes {
validation::aggregate(&change)?;
let identity = match &change {
UsageAggregateChange::Put { aggregate } => &aggregate.identity,
UsageAggregateChange::Delete { identity } => identity,
};
let action = match &change {
UsageAggregateChange::Put { .. } => "usage.aggregate.put",
UsageAggregateChange::Delete { .. } => "usage.aggregate.delete",
};
let audit = audit_event(
actor,
action,
"usage_aggregate",
&identity.key,
vec!["aggregate".into()],
);
changes.push((change, audit));
}
self.store
.write_aggregates(StoredAggregateWrite { changes })
.await
}
async fn get_record(&self, request: GetUsageRecord) -> UsageDataResult<UsageRecord> {
validation::text(&request.id, 128)?;
self.store.get_record(request).await
}
async fn list_records(&self, request: ListUsageRecords) -> UsageDataResult<UsageRecordPage> {
validation::query(&request.filter, request.cursor.as_deref(), request.limit)?;
self.store.list_records(request).await
}
async fn write_records(
&self,
actor: &UsageActor,
request: WriteUsageRecords,
) -> UsageDataResult<UsageWriteResult> {
validation::bounded(&request, request.changes.len())?;
let mut changes = Vec::with_capacity(request.changes.len());
for change in request.changes {
validation::change(&change)?;
let action = match &change {
UsageRecordChange::Put { .. } => "usage.put",
UsageRecordChange::Update { .. } => "usage.update",
UsageRecordChange::Delete { .. } => "usage.delete",
};
let fields = match &change {
UsageRecordChange::Update { tokens, cost, .. } => tokens
.as_ref()
.map(|_| "tokens".to_owned())
.into_iter()
.chain(cost.as_ref().map(|_| "cost".to_owned()))
.collect(),
_ => vec!["record".into()],
};
let audit = audit_event(actor, action, "model_request", change.id(), fields);
changes.push((change, audit));
}
self.store.write_records(StoredUsageWrite { changes }).await
}
async fn evaluate_pricing(
&self,
request: EvaluateUsagePricing,
) -> UsageDataResult<UsagePricingEvaluation> {
validation::bounded(&request, request.items.len())?;
let stored = self
.settings
.load_pricing()
.await
.map_err(|_| UsageDataError::unavailable("读取价目失败"))?;
let prices = gateway_core::metering::merge_pricing(stored.synced, &stored.overrides);
let mut items = Vec::with_capacity(request.items.len());
for input in request.items {
validation::text(&input.id, 128)?;
items.push(self.evaluate_one(input, &prices)?);
}
Ok(UsagePricingEvaluation { items })
}
}
fn audit_event(
actor: &UsageActor,
action: &str,
entity_kind: &str,
id: &str,
changed_fields: Vec<String>,
) -> AdminAuditEvent {
let actor_ref = match actor {
UsageActor::Plugin {
instance_id,
publisher,
name,
..
} => format!("plugin:{publisher}/{name}:{instance_id}"),
UsageActor::Admin { subject } => format!("admin:{subject}"),
UsageActor::System => "system".into(),
};
AdminAuditEvent {
id: format!("audit_{}", uuid::Uuid::now_v7().simple()),
actor_kind: AuditActorKind::System,
actor_admin_user_id: None,
actor_ref,
request_id: None,
action: action.into(),
entity_kind: entity_kind.into(),
entity_ref: id.to_owned(),
config_revision: None,
changed_fields,
occurred_at: chrono::Utc::now(),
}
}
@@ -0,0 +1,87 @@
//! 复用 Provider 的单请求 Token 计价规则,不写入历史或账本
use super::{DefaultUsageDataService, validation};
use crate::{
model::usage_data::*,
ports::provider::{ProviderAdminErrorKind, ProviderUsagePricingInput},
};
use gateway_core::{
metering::{PricingOverrides, Usage},
routing::ProviderKind,
};
impl DefaultUsageDataService {
pub(super) fn evaluate_one(
&self,
input: UsagePricingInput,
prices: &PricingOverrides,
) -> UsageDataResult<UsagePricingResult> {
validation::text(&input.provider, 128)?;
validation::text(&input.model, 128)?;
validation::tokens(&input.tokens)?;
let unavailable = |reason: &str| UsagePricingResult::NotEstimable {
id: input.id.clone(),
reason: reason.into(),
};
let kind = ProviderKind::new(input.provider.clone())
.map_err(|_| UsageDataError::invalid("Provider 无效"))?;
let Ok(provider) = self.providers.require(&kind) else {
return Ok(unavailable("Provider 没有可用的计价规则"));
};
let tokens = &input.tokens;
let price = match &input.price {
UsagePriceSelection::Current => prices
.get(&input.provider)
.and_then(|models| models.get(&input.model)),
UsagePriceSelection::Override { price } => {
price.validate().map_err(UsageDataError::invalid)?;
Some(price.as_ref())
}
};
let request = ProviderUsagePricingInput {
upstream_model_id: input.model.clone(),
service_tier: input.service_tier.clone(),
usage: Usage {
input_tokens: tokens.input_tokens.map(UsageCount::get),
output_tokens: tokens.output_tokens.map(UsageCount::get),
cached_input_tokens: tokens.cached_input_tokens.map(UsageCount::get),
cache_write_input_tokens: tokens.cache_write_input_tokens.map(UsageCount::get),
reasoning_output_tokens: tokens.reasoning_output_tokens.map(UsageCount::get),
image_input_tokens: tokens.image_input_tokens.map(UsageCount::get),
image_output_tokens: tokens.image_output_tokens.map(UsageCount::get),
total_tokens: tokens.total_tokens.map(UsageCount::get),
},
};
let calculated = match provider.evaluate_usage_pricing(&request, price) {
Ok(Some(value)) => value,
Ok(None) => return Ok(unavailable("用量不完整、包含图像或价格不可用")),
Err(error)
if matches!(
error.kind(),
ProviderAdminErrorKind::Unsupported | ProviderAdminErrorKind::Invalid
) =>
{
return Ok(unavailable("Provider 不支持该计价输入"));
}
Err(_) => return Err(UsageDataError::unavailable("Provider 计价失败")),
};
Ok(UsagePricingResult::Estimated {
id: input.id,
cost: UsageCost::Known {
source: UsageCostSource::Calculated,
amount: calculated.total_amount().amount().canonical(),
currency: calculated.total_amount().currency().as_str().into(),
},
breakdown: UsagePricingBreakdown {
input_amount: calculated.input_amount().amount().canonical(),
output_amount: calculated.output_amount().amount().canonical(),
cache_read_amount: calculated.cache_read_amount().amount().canonical(),
cache_write_amount: calculated.cache_write_amount().amount().canonical(),
standard_amount: calculated.standard_amount().amount().canonical(),
long_context_billing_applied: calculated.long_context_billing_applied(),
service_tier: calculated.service_tier().map(str::to_owned),
custom_multiplier_bps: calculated.custom_multiplier_bps(),
},
})
}
}
@@ -0,0 +1,194 @@
//! 用量输入沿用原生记录的数值、时间和资源边界
use crate::model::usage_data::*;
use gateway_core::metering::Decimal;
use serde::Serialize;
pub(super) fn text(value: &str, max: usize) -> UsageDataResult<()> {
if value.trim().is_empty() || value.len() > max || value.chars().any(char::is_control) {
return Err(UsageDataError::invalid("文本为空、过长或包含控制字符"));
}
Ok(())
}
pub(super) fn bounded(value: &impl Serialize, count: usize) -> UsageDataResult<()> {
if count == 0
|| count > USAGE_DATA_MAX_TARGETS
|| serde_json::to_vec(value)
.map_err(|_| UsageDataError::invalid("请求编码失败"))?
.len()
> USAGE_DATA_MAX_REQUEST_BYTES
{
return Err(UsageDataError::new(
UsageErrorKind::CapacityExceeded,
"每次操作接受 1~500 项且不超过 2 MiB",
));
}
Ok(())
}
pub(super) fn tokens(value: &UsageTokens) -> UsageDataResult<()> {
if [
value.input_tokens,
value.output_tokens,
value.cached_input_tokens,
value.cache_write_input_tokens,
value.reasoning_output_tokens,
value.image_input_tokens,
value.image_output_tokens,
value.total_tokens,
]
.into_iter()
.flatten()
.any(|v| v.get() > i64::MAX as u64)
{
return Err(UsageDataError::invalid("Token 超出原生存储范围"));
}
Ok(())
}
fn cost(value: &UsageCost) -> UsageDataResult<()> {
if let UsageCost::Known {
amount, currency, ..
} = value
{
let _ = amount
.parse::<Decimal>()
.map_err(|_| UsageDataError::invalid("费用不是有效十进制数"))?;
if currency.len() != 3 || !currency.bytes().all(|b| b.is_ascii_uppercase()) {
return Err(UsageDataError::invalid("费用或币种不符合原生存储范围"));
}
}
Ok(())
}
pub(super) fn change(value: &UsageRecordChange) -> UsageDataResult<()> {
text(value.id(), 128)?;
match value {
UsageRecordChange::Put { data, .. } => {
text(&data.client_key_ref, 128)?;
for v in [
&data.provider,
&data.account_ref,
&data.requested_model,
&data.model,
&data.service_tier,
]
.into_iter()
.flatten()
{
text(v, 128)?;
}
for v in [
&data.protocol,
&data.operation,
&data.endpoint,
&data.client_transport,
] {
text(v, 256)?;
}
if data.outcome == UsageOutcome::Running
|| data
.completed_at_ms
.is_none_or(|end| end < data.started_at_ms)
{
return Err(UsageDataError::invalid("只允许写入已结束的请求明细"));
}
for at in [Some(data.started_at_ms), data.completed_at_ms]
.into_iter()
.flatten()
{
if chrono::DateTime::from_timestamp_millis(at).is_none() {
return Err(UsageDataError::invalid("请求时间无效"));
}
}
tokens(&data.tokens)?;
cost(&data.cost)?;
}
UsageRecordChange::Update {
tokens: new_tokens,
cost: new_cost,
..
} => {
if new_tokens.is_none() && new_cost.is_none() {
return Err(UsageDataError::invalid("修改必须指定 Token 或费用"));
}
if let Some(value) = new_tokens {
tokens(value)?;
}
if let Some(value) = new_cost {
cost(value)?;
}
}
UsageRecordChange::Delete { .. } => {}
}
Ok(())
}
pub(super) fn query(
filter: &UsageRecordFilter,
cursor: Option<&str>,
limit: Option<u16>,
) -> UsageDataResult<()> {
if limit.is_some_and(|v| v == 0 || v > USAGE_DATA_MAX_PAGE_SIZE)
|| filter.start_ms >= filter.end_ms
|| chrono::DateTime::from_timestamp_millis(filter.start_ms).is_none()
|| chrono::DateTime::from_timestamp_millis(filter.end_ms).is_none()
{
return Err(UsageDataError::invalid("查询时间范围或分页大小无效"));
}
for v in [
&filter.provider,
&filter.model,
&filter.account_ref,
&filter.client_key_ref,
]
.into_iter()
.flatten()
{
text(v, 128)?;
}
if cursor.is_some_and(|v| v.len() > 4096) {
return Err(UsageDataError::invalid("游标过长"));
}
Ok(())
}
pub(super) fn aggregate(value: &UsageAggregateChange) -> UsageDataResult<()> {
let identity = match value {
UsageAggregateChange::Put { aggregate } => &aggregate.identity,
UsageAggregateChange::Delete { identity } => identity,
};
text(&identity.key, 128)?;
if chrono::DateTime::from_timestamp_millis(identity.bucket_start_ms).is_none() {
return Err(UsageDataError::invalid("聚合归属时间无效"));
}
let d = &identity.dimensions;
for value in [
&d.provider,
&d.model,
&d.account_ref,
&d.client_key_ref,
&d.operation,
&d.transport,
]
.into_iter()
.flatten()
{
text(value, 128)?;
}
if let UsageAggregateChange::Put { aggregate } = value {
if aggregate.costs.len() > 32 {
return Err(UsageDataError::invalid("费用币种数量超过限制"));
}
for (currency, amount) in &aggregate.costs {
if currency.len() != 3
|| !currency.bytes().all(|b| b.is_ascii_uppercase())
|| amount.len() > 128
{
return Err(UsageDataError::invalid("聚合费用格式无效"));
}
amount
.strip_prefix('-')
.unwrap_or(amount)
.parse::<Decimal>()
.map_err(|_| UsageDataError::invalid("聚合费用不是有效十进制数"))?;
}
}
Ok(())
}
@@ -4,7 +4,13 @@ use gateway_admin::model::retention::{RetentionPolicy, RetentionTarget};
#[test]
fn retention_policy_rejects_invalid_windows_and_preserves_independent_windows() {
for (usage, ops, audit) in [(30, 1, 1), (31, 0, 1), (31, 1, 0)] {
for (usage, ops, audit) in [
(0, 1, 1),
(32, 1, 1),
(u32::MAX, 1, 1),
(1, 0, 1),
(1, 1, 0),
] {
assert!(RetentionPolicy::try_new(usage, ops, audit).is_err());
}
let policy = RetentionPolicy::try_new(31, 7, 90).unwrap();
@@ -12,3 +18,13 @@ fn retention_policy_rejects_invalid_windows_and_preserves_independent_windows()
assert_eq!(policy.days(RetentionTarget::OpsEvents), 7);
assert_eq!(policy.days(RetentionTarget::AdminAuditEvents), 90);
}
#[test]
fn request_retention_accepts_short_windows_without_changing_other_targets() {
for days in [1, 3, 7, 15, 30, 31] {
let policy = RetentionPolicy::try_new(days, 30, 90).unwrap();
assert_eq!(policy.days(RetentionTarget::ModelRequests), days);
assert_eq!(policy.days(RetentionTarget::OpsEvents), 30);
assert_eq!(policy.days(RetentionTarget::AdminAuditEvents), 90);
}
}
@@ -317,6 +317,7 @@ impl AccountRuntimeStore for FakeRuntimeStore {
fn group_record() -> AccountGroupRecord {
let now = Utc::now();
AccountGroupRecord {
allow_image_generation: true,
fast_mode: FastMode::Default,
id: group_id(),
name: "Primary".to_owned(),
@@ -390,6 +391,7 @@ async fn invalid_group_fields_are_rejected_before_store_mutations() {
.create(
&context,
CreateAccountGroup {
allow_image_generation: true,
fast_mode: FastMode::Default,
name: " invalid".to_owned(),
description: None,
@@ -403,6 +405,7 @@ async fn invalid_group_fields_are_rejected_before_store_mutations() {
.update(
&context,
UpdateAccountGroup {
allow_image_generation: None,
fast_mode: None,
id: record.id,
name: "valid".to_owned(),
@@ -15,6 +15,7 @@ mod plugins;
mod proxies;
mod settings;
mod system;
mod usage_data;
use std::{
str::FromStr,
@@ -273,6 +274,7 @@ impl AdminHarness {
self.plugin_store,
Arc::new(plugins::TestPluginPorts),
Arc::new(plugins::TestPluginPorts),
Arc::new(plugins::TestPluginPorts),
),
gateway_admin::AdminRuntimePorts {
timezone: Default::default(),
@@ -747,6 +749,14 @@ impl ClientKeyStore for UnavailableStore {
#[async_trait]
impl ObservabilityStore for UnavailableStore {
async fn diagnostic_errors(
&self,
_target: &gateway_admin::model::diagnostics::DiagnosticTarget,
_limit: u16,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::OpsError>> {
Err(unavailable("diagnostic errors"))
}
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
@@ -945,6 +945,14 @@ impl FixtureObservabilityStore {
#[async_trait]
impl ObservabilityStore for FixtureObservabilityStore {
async fn diagnostic_errors(
&self,
_target: &gateway_admin::model::diagnostics::DiagnosticTarget,
_limit: u16,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::OpsError>> {
Err(super::unavailable("diagnostic errors"))
}
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
@@ -1194,6 +1202,7 @@ fn observation_range(end: DateTime<Utc>) -> TimeRange {
fn usage_query(now: DateTime<Utc>) -> UsageQuery {
UsageQuery {
status: Default::default(),
range: observation_range(now),
filter: UsageFilter::default(),
current_page: 1,
@@ -1208,6 +1217,7 @@ fn total_record(
now: DateTime<Utc>,
) -> UsageListRecord {
UsageListRecord {
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
client_api_key_name: Some("Production".to_owned()),
id: id.to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -1,5 +1,7 @@
//! 插件管理用例的测试入口与共用端口替身
mod usage_data;
use async_trait::async_trait;
use gateway_admin::model::plugins::distribution::{
DownloadedPlugin, GithubReleaseQuery, PluginRelease, RemotePluginLocation, SourceCredential,
@@ -0,0 +1,24 @@
//! 通用用量端口的插件管理测试夹具,非本测试调用路径禁止访问
use super::TestPluginPorts;
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
#[async_trait::async_trait]
impl UsageDataStore for TestPluginPorts {
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
unreachable!("unexpected usage call")
}
async fn list_records(&self, _: ListUsageRecords) -> UsageDataResult<UsageRecordPage> {
unreachable!("unexpected usage call")
}
async fn write_records(&self, _: StoredUsageWrite) -> UsageDataResult<UsageWriteResult> {
unreachable!("unexpected usage call")
}
async fn list_aggregates(&self, _: ListUsageAggregates) -> UsageDataResult<UsageAggregatePage> {
unreachable!("unexpected usage call")
}
async fn write_aggregates(
&self,
_: StoredAggregateWrite,
) -> UsageDataResult<UsageAggregateWriteResult> {
unreachable!("unexpected usage call")
}
}
@@ -0,0 +1,119 @@
//! 用量服务的 wire 精度、可信主体和独立 SDK 合同校验
use gateway_admin::model::usage_data::*;
use serde_json::json;
#[test]
fn counters_preserve_integer_precision_and_reject_json_numbers() {
assert_eq!(
serde_json::from_value::<UsageCount>(json!(u64::MAX.to_string()))
.unwrap()
.get(),
u64::MAX
);
assert_eq!(
serde_json::from_value::<UsageDelta>(json!(i64::MIN.to_string()))
.unwrap()
.get(),
i64::MIN
);
for value in [json!(1), json!("1.5"), json!("+1")] {
assert!(serde_json::from_value::<UsageCount>(value.clone()).is_err());
assert!(serde_json::from_value::<UsageDelta>(value).is_err());
}
}
#[tokio::test]
async fn usage_mutations_require_trusted_actor_and_old_workflow_operations_are_absent() {
let admin = super::AdminHarness::new().build().await;
let registry = admin.public_services();
let denied = registry
.call("usage.records.write", json!({"changes":[]}))
.await
.unwrap_err();
assert!(denied.message.contains("trusted caller"));
assert!(
registry
.call(
"usage.records.write",
json!({"changes":[],"actor":{"kind":"system"}})
)
.await
.is_err()
);
let ops = registry
.operations()
.filter(|op| op.starts_with("usage.") || *op == "pricing.evaluate")
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
ops,
std::collections::BTreeSet::from([
"usage.records.get",
"usage.records.list",
"usage.records.write",
"usage.aggregates.list",
"usage.aggregates.write",
"pricing.evaluate"
])
);
}
#[test]
fn sdk_wire_types_match_host_contracts() {
use quote::ToTokens as _;
let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
let items = |path: std::path::PathBuf| {
syn::parse_file(&std::fs::read_to_string(path).unwrap())
.unwrap()
.items
.into_iter()
.filter(|item| {
matches!(
item,
syn::Item::Struct(_)
| syn::Item::Enum(_)
| syn::Item::Impl(_)
| syn::Item::Type(_)
)
})
.map(|item| item.into_token_stream().to_string())
.collect::<Vec<_>>()
};
for name in [
"value.rs",
"record.rs",
"query.rs",
"aggregate.rs",
"pricing.rs",
"error.rs",
] {
assert_eq!(
items(root.join("src/model/usage_data").join(name)),
items(
root.join("../gateway-plugin/sdk/src/call/services/usage")
.join(name)
),
"{name}"
);
}
}
#[tokio::test]
async fn batch_and_wire_limits_fail_before_reaching_storage() {
let admin = super::AdminHarness::new().build().await;
let registry = admin.public_services();
let changes = (0..501)
.map(|i| json!({"kind":"delete","id":format!("item_{i}")}))
.collect::<Vec<_>>();
let too_many = gateway_admin::public_service::scope_actor(
UsageActor::System,
registry.call("usage.records.write", json!({"changes":changes})),
)
.await
.unwrap_err();
assert_eq!(too_many.kind, "capacity_exceeded");
let oversized = registry
.call(
"usage.records.get",
json!({"id":"x".repeat(USAGE_DATA_MAX_REQUEST_BYTES)}),
)
.await
.unwrap_err();
assert_eq!(oversized.kind, "capacity_exceeded");
}
+1
View File
@@ -25,6 +25,7 @@ gateway-admin = { path = "../gateway-admin" }
gateway-core = { path = "../gateway-core" }
gateway-protocol = { path = "../gateway-protocol" }
headers = "0.4.1"
regex = "1.13.1"
secrecy = "0.10"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
@@ -91,9 +91,15 @@ impl ListAccountGroupsQuery {
}
}
const fn default_allow_image_generation() -> bool {
true
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateAccountGroupRequest {
#[serde(default = "default_allow_image_generation")]
allow_image_generation: bool,
#[serde(default)]
fast_mode: FastMode,
name: String,
@@ -104,6 +110,7 @@ struct CreateAccountGroupRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct UpdateAccountGroupRequest {
allow_image_generation: Option<bool>,
fast_mode: Option<FastMode>,
id: String,
name: String,
@@ -120,6 +127,7 @@ struct AccountGroupIdRequest {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupView {
allow_image_generation: bool,
fast_mode: FastMode,
id: String,
name: String,
@@ -167,6 +175,7 @@ impl From<(AccountGroupRecord, crate::time::TimePresenter)> for AccountGroupView
name: record.name,
description: record.description,
fast_mode: record.fast_mode,
allow_image_generation: record.allow_image_generation,
color: record.color.as_str().to_owned(),
enabled: record.enabled,
member_count: record.member_count,
@@ -362,6 +371,7 @@ where
.create(
&auth.context().mutation_context(),
CreateAccountGroup {
allow_image_generation: request.allow_image_generation,
name: request.name,
description: request.description,
fast_mode: request.fast_mode,
@@ -390,6 +400,7 @@ where
.update(
&auth.context().mutation_context(),
UpdateAccountGroup {
allow_image_generation: request.allow_image_generation,
id: group_id(request.id)?,
name: request.name,
description: request.description,
@@ -0,0 +1,136 @@
//! 诊断文件的 HTTP 下载与受控数据投影,普通详情响应不携带时间线
use axum::{extract::State, http::header, response::IntoResponse};
use gateway_admin::model::{
diagnostics::{DiagnosticBundle, DiagnosticTarget},
observability::OpsError,
};
use serde::Deserialize;
use serde_json::{Value, json};
use super::{
AdminAuth, AdminError, AdminQuery, map_service_error, ops_error_view, usage_detail_view,
};
use crate::{auth::SessionState, time::TimePresenter};
mod privacy;
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(super) struct ExportQuery {
request_id: Option<String>,
error_id: Option<String>,
}
pub(super) async fn download<S>(
_auth: AdminAuth,
State(state): State<S>,
AdminQuery(query): AdminQuery<ExportQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: SessionState + Send + Sync,
{
let target = match (query.request_id, query.error_id) {
(Some(id), None) => DiagnosticTarget::Request(id),
(None, Some(id)) => DiagnosticTarget::Operation(id),
_ => {
return Err(map_service_error(
gateway_admin::model::AdminError::invalid("请指定一个请求 ID 或错误记录 ID"),
));
}
};
let bundle = state
.admin_services()
.observability()
.export_diagnostics(target)
.await
.map_err(map_service_error)?;
let time = TimePresenter::new(state.admin_services().timezone());
let document = privacy::redact(document(bundle, time));
let body = serde_json::to_vec_pretty(&document).map_err(|_| {
map_service_error(gateway_admin::model::AdminError::internal(
"诊断文件生成失败",
))
})?;
Ok((
[
(header::CONTENT_TYPE, "application/json; charset=utf-8"),
(
header::CONTENT_DISPOSITION,
"attachment; filename=diagnostics.json",
),
(header::CACHE_CONTROL, "private, no-store"),
(header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
],
body,
))
}
fn document(bundle: DiagnosticBundle, time: TimePresenter) -> Value {
let target = match &bundle.target {
DiagnosticTarget::Request(id) => json!({"requestId": id}),
DiagnosticTarget::Operation(id) => json!({"errorId": id}),
};
let requests = bundle.requests.into_iter().map(|record| {
let evidence = record.detail.map(|mut detail| {
let trace = detail.trace.take();
let related = std::mem::take(&mut detail.related_requests);
let configuration = json!({
"revision": detail.request.config_revision,
"routingScope": detail.request.routing_scope,
"routingGroupIds": detail.request.routing_group_refs,
"deadlineAt": detail.request.deadline_at,
});
// 页面 DTO 没有展示的尝试原因与等待事实仍属于导出证据
let attempt_details = detail.attempts.iter().map(|attempt| json!({
"id": attempt.id, "attemptIndex": attempt.attempt_index,
"component": attempt.component, "operation": attempt.operation,
"message": attempt.message, "retryAfterMs": attempt.retry_after_ms,
})).collect::<Vec<_>>();
json!({
"detail": usage_detail_view(detail, time), "trace": trace,
"relatedRequests": related, "incidentConfiguration": configuration,
"attemptDetails": attempt_details,
})
});
json!({
"requestId": record.request_id,
"evidence": evidence,
"errors": record.errors.into_iter().map(|error| export_error(error, time)).collect::<Vec<_>>(),
})
}).collect::<Vec<_>>();
json!({
"schemaVersion": 5,
"source": "gateway.diagnostics",
"exportPolicy": "support_credentials_redacted",
"target": target,
"collectedAt": bundle.collected_at,
"environment": {
"scope": "export_time", "observedAt": bundle.collected_at,
"gateway": bundle.environment.map(|version| json!({
"version": version.version, "gitSha": version.git_sha,
"buildTime": version.build_time, "deploymentMode": version.deployment_mode,
})),
},
"error": bundle.error.map(|error| export_error(error, time)),
"requests": requests,
"availability": {"gaps": bundle.gaps.into_iter().map(|gap| json!({"section": gap.section, "reason": gap.reason})).collect::<Vec<_>>()},
"omitted": ["credentials_and_authentication", "account_contact_details", "uncollected_wire_payloads", "unrelated_requests_and_system_secrets"],
"notes": [
"诊断包包含错误原因链、上游错误与业务上下文,用于故障分析",
"认证字段、账号联系信息与可识别的内嵌凭据已脱敏,任意业务内容不属于自动识别范围",
"environment 为导出时环境,incidentConfiguration 与 request.routed 为事发配置事实,当前版本不能代替事发版本",
"缺失、截断、淘汰与读取失败见 availability.gaps,没有记录不表示没有发生",
"attempts 为异步 best-effort 投影,firstTokenMs 为空不证明没有输出",
"未采集的请求正文、过期记录、上游内部状态与退出前未落盘信息无法由导出补回"
],
})
}
// 复用错误字段投影,原始详情直接进入脱敏遍历,避免先为页面缩进再解析一遍
fn export_error(mut error: OpsError, time: TimePresenter) -> Value {
let details = error.error_details.take();
let mut value = json!(ops_error_view(error, time));
value["errorDetails"] = json!(details);
value
}
@@ -0,0 +1,220 @@
//! 诊断文件按认证与观测合同移除敏感字段,保留错误事实并明确导出截断
use regex::Regex;
use serde::Serialize;
use serde_json::{Value, json};
use std::sync::LazyLock;
const REDACTED: &str = "<redacted>";
const MAX_TEXT: usize = 256 * 1024;
static INLINE: LazyLock<Vec<(Regex, &str)>> = LazyLock::new(|| {
[
(r"(?i)\b((?:set-cookie|cookie2?|(?:proxy-)?authorization|x-api-key|x-oai-is(?:-update)?|x-oai-attestation)\s*:\s*)[^\r\n]+", "${1}<redacted>"),
(r#"(?i)\b(Bearer|Basic)\s+[^\s"'<>;,]+"#, "${1} <redacted>"),
// 官方 TokenData 的 access_token / id_token 为 JWT,不按猜测的令牌前缀过滤普通文本
(r"\beyJ[\w-]+\.[\w-]+\.[\w-]+\b", "<redacted>"),
(r"(?i)([a-z][a-z0-9+.-]*://)[^\s/@]+@", "${1}<redacted>@"),
(r#"(?i)\b((?:(?:personal[_-]?)?access[_-]?token|refresh[_-]?token|id[_-]?token|(?:(?:openai|x)[_-]?)?api[_-]?key|(?:current[_-]?|new[_-]?)?password|authorization|proxy[_-]?authorization|x[_-]?oai[_-]?(?:is(?:[_-]?update)?|attestation))["']?\s*[=:]\s*)("[^"\r\n]*"|'[^'\r\n]*'|[^\s&,;]+)"#, "${1}<redacted>"),
].into_iter().map(|(pattern, replacement)| (Regex::new(pattern).expect("static credential pattern"), replacement)).collect()
});
#[derive(Default, Serialize)]
#[serde(rename_all = "camelCase")]
struct Report {
credential_fields: usize,
personal_fields: usize,
inline_credentials: usize,
truncated_values: usize,
}
struct Redactor {
report: Report,
bytes: usize,
nodes: usize,
}
pub(super) fn redact(mut value: Value) -> Value {
let mut redactor = Redactor {
report: Report::default(),
bytes: 8 * 1024 * 1024,
nodes: 100_000,
};
// 先处理有界文件信息,正文耗尽预算也不能删除 schema、目标或既有缺口
let mut error = value["error"].take();
let mut requests = value["requests"].take();
redactor.visit(&mut value, 0);
redactor.visit(&mut error, 0);
redactor.visit(&mut requests, 0);
value["error"] = error;
value["requests"] = requests;
if redactor.report.truncated_values > 0 {
// 缺口清单可能也触及预算,顶层统计始终保留以免截断变成静默丢失
value["exportTruncated"] = json!(true);
if !value["availability"]["gaps"].is_array() {
value["availability"] = json!({"gaps": []});
}
value["availability"]["gaps"]
.as_array_mut()
.expect("gap array")
.push(json!({
"section": "export", "reason": "export_size_limit",
}));
}
value["redaction"] = json!(redactor.report);
value
}
fn normalized(key: &str) -> String {
key.chars()
.filter(|character| !matches!(character, '_' | '-'))
.flat_map(char::to_lowercase)
.collect()
}
// OAuth 字段来自 Provider credential schema 与官方 Codex TokenData;API Key 来自
// CodexCprApiKeyExportAccount / AuthDotJson;密码来自管理登录合同,头部来自传输身份合同
// 不用 token、secret、credential 等泛称推断敏感性,避免删除业务错误上下文
fn secret(key: &str) -> bool {
matches!(
key,
"accesstoken"
| "personalaccesstoken"
| "refreshtoken"
| "idtoken"
| "apikey"
| "openaiapikey"
| "password"
| "currentpassword"
| "newpassword"
| "authorization"
| "proxyauthorization"
| "xapikey"
| "cookie"
| "cookie2"
| "setcookie"
| "xoaiis"
| "xoaiisupdate"
| "xoaiattestation"
)
}
impl Redactor {
fn omitted(&mut self) -> Value {
self.report.truncated_values += 1;
json!({"omitted": "export_size_limit"})
}
fn text(&mut self, text: &str) -> String {
let mut result = text.to_owned();
for (pattern, replacement) in INLINE.iter() {
self.report.inline_credentials += pattern.find_iter(&result).count();
result = pattern.replace_all(&result, *replacement).into_owned();
}
result
}
fn visit(&mut self, value: &mut Value, depth: usize) {
if self.nodes == 0 || self.bytes == 0 || depth > 32 {
*value = self.omitted();
return;
}
self.nodes -= 1;
// 预算同时计入缩进和容器分隔符,避免深层 JSON 的下载体积显著超过正文预算
self.bytes = self.bytes.saturating_sub(24 + depth * 4);
match value {
Value::String(text) => {
if text.trim_start().starts_with('{') || text.trim_start().starts_with('[') {
if text.len() > MAX_TEXT {
*value = self.omitted();
return;
}
if let Ok(parsed) = serde_json::from_str::<Value>(text) {
*value = parsed;
self.visit(value, depth + 1);
return;
}
}
let cleaned = self.text(text);
let limit = MAX_TEXT.min(self.bytes / 2);
if cleaned.len() > limit {
let boundary = cleaned.floor_char_boundary(limit);
*value = json!({"text": &cleaned[..boundary], "omittedBytes": cleaned.len() - boundary});
self.report.truncated_values += 1;
} else {
*text = cleaned;
}
self.bytes = self.bytes.saturating_sub(value.to_string().len());
}
Value::Array(items) => {
// 头部也可能以二元组承载,不能只识别对象字段名
if items.len() == 2
&& items[0]
.as_str()
.is_some_and(|name| secret(&normalized(name)))
{
items[1] = json!(REDACTED);
self.report.credential_fields += 1;
}
let original_len = items.len();
let mut kept = 0;
for item in items.iter_mut().take(4096) {
if self.nodes == 0 || self.bytes == 0 {
break;
}
self.visit(item, depth + 1);
kept += 1;
}
if kept < original_len {
items.truncate(kept);
items.push(json!({"omittedItems": original_len - kept}));
self.report.truncated_values += 1;
}
}
Value::Object(object) => {
if object
.get("name")
.and_then(Value::as_str)
.is_some_and(|name| secret(&normalized(name)))
&& object.contains_key("value")
{
object.insert("value".to_owned(), json!(REDACTED));
self.report.credential_fields += 1;
}
let fields = std::mem::take(object);
let total = fields.len();
for (key, mut item) in fields {
if object.len() >= 4096 || self.nodes == 0 || self.bytes == 0 {
break;
}
let normalized = normalized(&key);
if secret(&normalized) {
item = json!(REDACTED);
self.report.credential_fields += 1;
} else if matches!(
normalized.as_str(),
"email"
| "accountemail"
| "clientip"
| "accountname"
| "accountnotes"
| "credentialname"
| "clientapikeyname"
| "accountlabel"
) {
item = json!(REDACTED);
self.report.personal_fields += 1;
} else {
self.visit(&mut item, depth + 1);
}
self.bytes = self.bytes.saturating_sub(key.len());
object.insert(self.text(&key), item);
}
if object.len() < total {
object.insert("_omittedFields".to_owned(), json!(total - object.len()));
self.report.truncated_values += 1;
}
}
_ => {}
}
}
}
@@ -16,6 +16,7 @@ use super::{
wire::map_admin_service_error,
};
mod export;
mod presenter;
mod query;
mod routes;
@@ -262,6 +262,7 @@ pub(crate) fn usage_list_record_view(
.clone()
.or_else(|| record.requested_model_id.clone());
UsageListRecordView {
logical_outcome: record.outcome.as_str().to_owned(),
client_api_key_name: record.client_api_key_name,
id: record.id,
provider: record.provider_kind,
@@ -535,8 +536,6 @@ pub(crate) fn usage_detail_view(
time: crate::time::TimePresenter,
) -> UsageRecordDetailView {
UsageRecordDetailView {
trace: detail.trace,
related_requests: detail.related_requests,
request: usage_record_view(detail.request, time),
attempts: detail
.attempts
@@ -616,7 +615,7 @@ pub(crate) fn ops_error_view(
subagent_kind: error.subagent_kind,
compact: error.compact,
message: error.message,
error_details: error.error_details,
error_details: error.error_details.map(error_details_view),
metadata: OpsErrorMetadataView {
source: error.source,
component: error.component,
@@ -640,6 +639,20 @@ pub(crate) fn ops_error_view(
}
}
// 受控错误正文保持原文语义,嵌套 JSON 只在展示投影中展开缩进
fn error_details_view(raw: String) -> String {
let Ok(mut value) = serde_json::from_str::<Value>(&raw) else {
return raw;
};
if let Some(upstream) = value.get_mut("upstream")
&& let Some(text) = upstream.as_str()
&& let Ok(parsed) = serde_json::from_str::<Value>(text)
{
*upstream = parsed;
}
serde_json::to_string_pretty(&value).unwrap_or(raw)
}
pub(crate) fn ops_page_view(
page: domain::OpsErrorPage,
time: crate::time::TimePresenter,
@@ -29,6 +29,7 @@ impl DashboardQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UsageQuery {
pub status: Option<String>,
pub current_page: Option<u32>,
pub page_size: Option<u16>,
pub kind: Option<String>,
@@ -330,6 +331,12 @@ pub(crate) fn usage_command(
let page_size_value =
DomainPageSize::new(page_size).map_err(|_| WireValidationError::new("pageSize"))?;
Ok(domain::UsageQuery {
status: match query.status.as_deref() {
None | Some("success") => domain::UsageRecordStatus::Success,
Some("error") => domain::UsageRecordStatus::Error,
Some("all") => domain::UsageRecordStatus::All,
Some(_) => return Err(WireValidationError::new("status")),
},
range: usage_range(
query.start_time.as_deref(),
query.end_time.as_deref(),
@@ -29,6 +29,10 @@ where
get(usage_insights_diagnostics::<S>),
)
.route("/api/admin/operations/errors", get(ops_errors::<S>))
.route(
"/api/admin/diagnostics/export",
get(super::export::download::<S>),
)
}
pub(crate) async fn dashboard_summary<S>(
@@ -86,6 +86,7 @@ pub struct ImageBillingView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageListRecordView {
pub logical_outcome: String,
pub client_api_key_name: Option<String>,
pub id: String,
pub provider: Option<String>,
@@ -275,8 +276,6 @@ pub struct UsageAttemptView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageRecordDetailView {
pub trace: Option<serde_json::Value>,
pub related_requests: Vec<serde_json::Value>,
#[serde(flatten)]
pub request: UsageRecordView,
pub attempts: Vec<UsageAttemptView>,
@@ -106,6 +106,7 @@ async fn create_route_should_create_an_empty_group() {
assert_eq!(value["data"]["record"]["memberCount"], 0);
assert_eq!(value["data"]["record"]["color"], "#A855F780");
assert_eq!(value["data"]["record"]["fastMode"], "default");
assert_eq!(value["data"]["record"]["allowImageGeneration"], true);
assert_eq!(value["data"]["configRevision"], 8);
}
@@ -393,3 +394,40 @@ async fn fast_mode_rejects_invalid_values_and_creates_groups_in_each_mode() {
}
}
}
#[tokio::test]
async fn image_generation_flag_round_trips_and_omitted_updates_preserve_it() {
let fixture = authenticated_fixture().await;
let created = request(router(&fixture), Method::POST, "/api/admin/account-groups/create", Some(json!({
"name":"No images", "description":null, "color":"#2563EBFF", "allowImageGeneration":false
})), true).await;
assert_eq!(created.status(), StatusCode::CREATED);
let created = response_json(created).await;
assert_eq!(created["data"]["record"]["allowImageGeneration"], false);
let id = created["data"]["id"].as_str().unwrap();
for (value, expected) in [
(None, false),
(Some(true), true),
(Some(false), false),
(None, false),
] {
let mut body =
json!({"id":id,"name":"Image policy","description":null,"color":"#2563EBFF"});
if let Some(value) = value {
body["allowImageGeneration"] = json!(value);
}
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/update",
Some(body),
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response_json(response).await["data"]["record"]["allowImageGeneration"],
expected
);
}
}
+30 -1
View File
@@ -185,6 +185,7 @@ impl AdminTestFixture {
plugin_ports.clone(),
plugin_ports.clone(),
plugin_ports.clone(),
plugin_ports.clone(),
);
let providers: Vec<Arc<dyn ProviderAdmin>> = vec![
Arc::new(UnusedProvider::new("openai", Arc::clone(&provider_error))),
@@ -710,6 +711,7 @@ impl MemoryAccountGroupStore {
(
primary_id.clone(),
AccountGroupRecord {
allow_image_generation: true,
fast_mode: FastMode::Default,
id: primary_id,
name: "Alpha routing".to_owned(),
@@ -732,6 +734,7 @@ impl MemoryAccountGroupStore {
(
secondary_id.clone(),
AccountGroupRecord {
allow_image_generation: true,
fast_mode: FastMode::Default,
id: secondary_id,
name: "Beta routing".to_owned(),
@@ -890,6 +893,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
let mut state = self.state.lock().expect("account groups");
let now = Utc::now();
let record = AccountGroupRecord {
allow_image_generation: command.allow_image_generation,
fast_mode: command.fast_mode,
id: command.id.clone(),
name: command.name,
@@ -925,6 +929,9 @@ impl AccountGroupStore for MemoryAccountGroupStore {
if let Some(fast_mode) = command.fast_mode {
record.fast_mode = fast_mode;
}
if let Some(allowed) = command.allow_image_generation {
record.allow_image_generation = allowed;
}
record.updated_at = Utc::now();
mutation(&mut state, command.id, true)
}
@@ -1380,6 +1387,27 @@ impl AccountRuntimeStore for UnusedStore {
#[async_trait]
impl ObservabilityStore for UnusedStore {
async fn diagnostic_errors(
&self,
target: &gateway_admin::model::diagnostics::DiagnosticTarget,
limit: u16,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::OpsError>> {
let errors = self.ops_errors.lock().expect("ops errors");
Ok(errors
.iter()
.filter(|error| match target {
gateway_admin::model::diagnostics::DiagnosticTarget::Request(id) => {
error.request_id.as_ref() == Some(id)
}
gateway_admin::model::diagnostics::DiagnosticTarget::Operation(id) => {
&error.event_id == id && error.source == "ops_event"
}
})
.take(usize::from(limit))
.cloned()
.collect())
}
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
@@ -1469,11 +1497,12 @@ impl ObservabilityStore for UnusedStore {
})
}
async fn usage_record_detail(&self, _: &str) -> AdminStoreResult<UsageDetail> {
async fn usage_record_detail(&self, request_id: &str) -> AdminStoreResult<UsageDetail> {
self.usage_detail
.lock()
.expect("usage detail")
.clone()
.filter(|detail| detail.request.id == request_id)
.ok_or_else(|| unavailable("usage detail"))
}
@@ -0,0 +1,402 @@
//! 诊断下载的权限、证据完整性、脱敏与部分失败合同
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header},
};
use chrono::Utc;
use gateway_admin::model::observability::UsageDetail;
use gateway_api::admin::observability;
use serde_json::{Value, json};
use tower::ServiceExt as _;
use super::response::{ops_error_record, usage_record_with_account};
use crate::admin::{AdminTestFixture, AdminTestState};
async fn export(
fixture: &AdminTestFixture,
query: &str,
authenticated: bool,
) -> (StatusCode, axum::http::HeaderMap, Value) {
let mut request = Request::builder()
.uri(format!("/api/admin/diagnostics/export{query}"))
.header("x-request-id", "diagnostics-export-test");
if authenticated {
request = request.header(header::COOKIE, "cpr_session=valid-session");
}
let response = observability::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(request.body(Body::empty()).unwrap())
.await
.unwrap();
let status = response.status();
let headers = response.headers().clone();
let body = to_bytes(response.into_body(), 20 * 1024 * 1024)
.await
.unwrap();
(status, headers, serde_json::from_slice(&body).unwrap())
}
#[tokio::test]
async fn diagnostic_export_requires_admin_and_exactly_one_target() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
assert_eq!(
export(&fixture, "?errorId=missing", false).await.0,
StatusCode::UNAUTHORIZED
);
for query in [
"",
"?requestId=a&errorId=b",
"?requestId=",
"?errorId=a&payload=untrusted",
] {
assert_eq!(
export(&fixture, query, true).await.0,
StatusCode::BAD_REQUEST,
"{query}"
);
}
assert_eq!(
export(&fixture, "?errorId=missing", true).await.0,
StatusCode::NOT_FOUND
);
}
#[tokio::test]
async fn diagnostic_export_keeps_retry_evidence_and_original_errors_but_not_credentials() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
error.error_details = Some(json!({
"causes": ["TLS certificate verification failed", "pool timeout after 5 seconds"],
"upstream": json!({"error": {"code": "unknown_vendor_failure", "message": "quota exhausted"},
"headers": {"Authorization": "fixture-authorization", "set-cookie": "fixture-cookie"},
"access_token": "fixture-access", "refreshToken": "fixture-refresh",
}).to_string(),
"endpoint": "postgres://fixture-user:fixture-password@database:5432/main",
"message": "connect failed: Bearer fixture-bearer, retry after 30 seconds",
}).to_string());
let mut unrelated = error.clone();
unrelated.request_id = Some("another-request".into());
unrelated.message = "unrelated evidence must not leak".into();
fixture
.ops_errors
.lock()
.unwrap()
.extend([error, unrelated]);
let trace = json!({"droppedEvents": 0, "events": [
{"stage": "request.routed", "data": {"accountPolicy": {"maxAccountRotations": 3}}},
{"stage": "account.selection", "data": {"candidates": [{"id": "acct_high", "weight": 2}, {"id": "acct_low", "weight": 1}]}},
{"stage": "attempt.failed", "data": {"kind": "unavailable", "recovery": {"replaySafe": true}}},
{"stage": "retry.decided", "data": {"phase": "execute", "retryable": false, "clientRetrySuppressed": true, "replayBarrierReached": false}},
]});
fixture.usage_detail.lock().unwrap().replace(UsageDetail {
request: usage_record_with_account(
"req_err",
"acct_err",
"private name",
"fixture@example.invalid",
"api_key",
Utc::now(),
),
attempts: Vec::new(),
trace: Some(trace.clone()),
related_requests: Vec::new(),
});
let (status, headers, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK, "{file}");
assert_eq!(
headers[header::CONTENT_DISPOSITION],
"attachment; filename=diagnostics.json"
);
assert_eq!(headers[header::CACHE_CONTROL], "private, no-store");
assert_eq!(file["schemaVersion"], 5);
assert_eq!(file["requests"][0]["evidence"]["trace"], trace);
assert_eq!(
file["requests"][0]["evidence"]["incidentConfiguration"]["revision"],
1
);
let details = &file["requests"][0]["errors"][0]["errorDetails"];
assert_eq!(
details["upstream"]["error"]["code"],
"unknown_vendor_failure"
);
assert_eq!(details["causes"][0], "TLS certificate verification failed");
assert_eq!(file["requests"][0]["errors"][0]["accountId"], "acct_err");
assert!(
file["availability"]["gaps"]
.as_array()
.unwrap()
.iter()
.any(|gap| gap["section"] == "environment")
);
let text = file.to_string();
for marker in [
"fixture-authorization",
"fixture-cookie",
"fixture-access",
"fixture-refresh",
"fixture-password",
"fixture-user",
"fixture-bearer",
"fixture@example.invalid",
"private name",
"unrelated evidence must not leak",
"203.0.113.8",
] {
assert!(!text.contains(marker), "leaked {marker}");
}
assert!(text.contains("retry after 30 seconds"));
}
#[tokio::test]
async fn background_error_exports_without_fabricating_a_model_request() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
error.source = "ops_event".into();
error.request_id = None;
error.component = "worker".into();
error.error_details =
Some(json!({"causes": ["database connection refused"], "upstream": null}).to_string());
fixture.ops_errors.lock().unwrap().push(error);
let (status, _, file) = export(&fixture, "?errorId=err_snapshot", true).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(file["requests"], json!([]));
assert_eq!(file["error"]["requestId"], Value::Null);
assert_eq!(
file["error"]["errorDetails"]["causes"][0],
"database connection refused"
);
}
#[tokio::test]
async fn error_survives_missing_request_and_nested_credentials_are_redacted() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
error.error_details = Some(json!({
"cookies": "Cookie: a=fixture-cookie-a; b=fixture-cookie-b\nupstream timeout",
"uri": "https://provider.invalid/path?api_key=fixture-query&model=gpt-test",
"headers": [["Authorization", "fixture-tuple"], {"name": "Cookie", "value": "fixture-header"}],
}).to_string());
fixture.ops_errors.lock().unwrap().push(error);
let (status, _, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK);
assert!(file["requests"][0]["evidence"].is_null());
assert!(
file["availability"]["gaps"]
.as_array()
.unwrap()
.iter()
.any(|gap| gap["reason"] == "request_detail_unavailable")
);
let text = file.to_string();
for marker in [
"fixture-cookie-a",
"fixture-cookie-b",
"fixture-query",
"fixture-tuple",
"fixture-header",
] {
assert!(!text.contains(marker), "leaked {marker}");
}
assert!(text.contains("upstream timeout"));
assert!(text.contains("model=gpt-test"));
}
#[tokio::test]
async fn oversized_error_export_reports_truncation_instead_of_silent_loss() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
error.source = "ops_event".into();
error.request_id = None;
error.message = "诊断文本".repeat(100_000);
fixture.ops_errors.lock().unwrap().push(error);
let (status, _, file) = export(&fixture, "?errorId=err_snapshot", true).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(file["exportTruncated"], true);
assert!(file["redaction"]["truncatedValues"].as_u64().unwrap() > 0);
assert!(file["error"]["message"]["omittedBytes"].as_u64().unwrap() > 0);
}
#[tokio::test]
async fn related_read_failures_and_limits_keep_primary_evidence() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut related = vec![json!({"requestId": "req_err", "relation": "recovers"})];
related.extend(
(0..10).map(
|index| json!({"requestId": format!("missing_{index}"), "relation": "recovered_by"}),
),
);
fixture.usage_detail.lock().unwrap().replace(UsageDetail {
request: usage_record_with_account(
"req_err",
"acct_err",
"fixture",
"fixture@example.invalid",
"api_key",
Utc::now(),
),
attempts: Vec::new(),
trace: None,
related_requests: related,
});
fixture
.ops_errors
.lock()
.unwrap()
.extend((0..130).map(|index| {
let mut error = ops_error_record();
error.event_id = format!("error_{index}");
error
}));
let (status, _, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(file["requests"].as_array().unwrap().len(), 8);
assert_eq!(file["requests"][0]["errors"].as_array().unwrap().len(), 128);
assert_eq!(
file["requests"][0]["evidence"]["detail"]["requestId"],
"req_err"
);
assert!(file["requests"][1]["evidence"].is_null());
let gaps = file["availability"]["gaps"].as_array().unwrap();
for reason in [
"related_request_limit",
"error_record_limit",
"request_detail_unavailable",
"not_recorded_or_retained",
] {
assert!(
gaps.iter().any(|gap| gap["reason"] == reason),
"missing {reason}"
);
}
}
#[tokio::test]
async fn credential_schemas_are_redacted_without_guessing_business_field_names() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
let business = json!({
"token": {"input": 128, "output": 64},
"credentials": {"status": "expired", "provider": "openai"},
"secret": "tool argument name", "clientSecret": "business field description",
"accountId": "account-correlation", "userId": "user-correlation",
"planType": "plus", "status": "active", "accessTokenExpiresAt": "2026-10-16T00:00:00Z",
"error": {"code": "unknown_code", "message": "upstream capacity exceeded"},
});
error.error_details = Some(json!({
"business": business,
"account": {"accessToken": "fixture-access", "refreshToken": "fixture-refresh", "idToken": "fixture-id"},
"auth": {"OPENAI_API_KEY": "fixture-key", "personal_access_token": "fixture-pat", "tokens": {
"access_token": "fixture-official-access", "refresh_token": "fixture-official-refresh",
"id_token": "fixture-official-id", "account_id": "account-correlation",
}},
}).to_string());
fixture.ops_errors.lock().unwrap().push(error);
let (status, _, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK);
let details = &file["requests"][0]["errors"][0]["errorDetails"];
assert_eq!(details["business"], business);
assert_eq!(
details["auth"]["tokens"]["account_id"],
"account-correlation"
);
assert!(!file.to_string().contains("fixture-"));
}
#[tokio::test]
async fn export_budget_keeps_format_target_and_existing_gap_records() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
fixture.usage_detail.lock().unwrap().replace(UsageDetail {
request: usage_record_with_account(
"req_err",
"acct_err",
"QA",
"qa@example.invalid",
"oauth",
Utc::now(),
),
attempts: Vec::new(),
trace: Some(json!({"events": vec![vec![true; 32]; 3500]})),
related_requests: Vec::new(),
});
let (status, _, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(file["schemaVersion"], 5);
assert_eq!(file["source"], "gateway.diagnostics");
assert_eq!(file["target"]["requestId"], "req_err");
assert_eq!(file["exportTruncated"], true);
let gaps = file["availability"]["gaps"].as_array().unwrap();
assert!(gaps.iter().any(|gap| gap["reason"] == "export_size_limit"));
assert!(
gaps.iter()
.any(|gap| gap["reason"] == "incident_policy_not_recorded")
);
assert!(serde_json::to_vec_pretty(&file).unwrap().len() < 16 * 1024 * 1024);
}
#[tokio::test]
async fn selected_event_is_exported_once_and_shares_the_total_error_budget() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut selected = ops_error_record();
selected.source = "ops_event".into();
let mut errors = vec![selected];
errors.extend((0..129).map(|index| {
let mut error = ops_error_record();
error.source = "ops_event".into();
error.event_id = format!("related_{index}");
error
}));
fixture.ops_errors.lock().unwrap().extend(errors);
let (status, _, file) = export(&fixture, "?errorId=err_snapshot", true).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(file["error"]["id"], "err_snapshot");
let errors = file["requests"][0]["errors"].as_array().unwrap();
assert_eq!(errors.len(), 127);
assert!(errors.iter().all(|error| error["id"] != "err_snapshot"));
assert!(
file["availability"]["gaps"]
.as_array()
.unwrap()
.iter()
.any(|gap| gap["reason"] == "error_record_limit")
);
}
#[tokio::test]
async fn inline_authentication_headers_are_redacted_without_losing_error_context() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut error = ops_error_record();
error.error_details = Some(json!({
"headers": "x-api-key: fixture-api-key\nx-oai-is: fixture-session\nx-oai-is-update: fixture-update\nx-oai-attestation: fixture-attestation fixture-attestation-tail\nauthorization: Custom fixture-auth\nproxy-authorization: Custom fixture-proxy\nupstream timeout",
"message": "access-token=fixture-access; refresh-token=fixture-refresh; model=gpt-test",
}).to_string());
fixture.ops_errors.lock().unwrap().push(error);
let (status, _, file) = export(&fixture, "?requestId=req_err", true).await;
assert_eq!(status, StatusCode::OK);
let text = file.to_string();
for marker in [
"fixture-api-key",
"fixture-session",
"fixture-update",
"fixture-attestation",
"fixture-attestation-tail",
"fixture-auth",
"fixture-proxy",
"fixture-access",
"fixture-refresh",
] {
assert!(!text.contains(marker), "leaked {marker}");
}
assert!(text.contains("upstream timeout"));
assert!(text.contains("model=gpt-test"));
}
@@ -1,4 +1,5 @@
//! Admin 观测 wire/路由的固定合同测试
mod export;
mod query;
mod response;
@@ -398,7 +398,7 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
}
#[tokio::test]
async fn zero_attempt_failure_detail_keeps_missing_upstream_facts_and_preparation_trace() {
async fn zero_attempt_failure_detail_omits_trace_and_keeps_missing_upstream_facts() {
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header},
@@ -470,7 +470,8 @@ async fn zero_attempt_failure_detail_keeps_missing_upstream_facts_and_preparatio
assert_eq!(data["attempts"], json!([]));
// 明细仍沿用全局 best-effort 合同,不从预备 trace 推断完整尝试列表
assert_eq!(data["attemptsComplete"], false);
assert_eq!(data["trace"], trace);
assert!(data.get("trace").is_none());
assert!(data.get("relatedRequests").is_none());
assert_eq!(data["logicalOutcome"], "failed");
assert_eq!(data["clientStatusCode"], 503);
assert_eq!(data["message"], "no upstream account is available");
@@ -495,8 +496,6 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
body::{Body, to_bytes},
http::{Request, StatusCode, header},
};
use chrono::Utc;
use gateway_admin::model::observability::OpsError;
use gateway_api::admin::observability;
use tower::ServiceExt as _;
@@ -508,63 +507,7 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
.ops_errors
.lock()
.expect("ops errors")
.push(OpsError {
client_api_key_name: Some("Production".to_owned()),
source: "model_request".to_owned(),
event_id: "err_snapshot".to_owned(),
request_id: Some("req_err".to_owned()),
attempt_index: Some(1),
client_api_key_ref: Some("key_err".to_owned()),
component: "model_request".to_owned(),
operation: "responses".to_owned(),
protocol: Some("openai".to_owned()),
client_transport: Some("http_sse".to_owned()),
requested_model_id: Some("gpt-5.4".to_owned()),
service_tier: Some("priority".to_owned()),
endpoint: Some("/v1/responses".to_owned()),
provider_kind: Some("openai".to_owned()),
provider_account_ref: Some("acct_err".to_owned()),
provider_account_name: None,
provider_account_email: Some("err@example.invalid".to_owned()),
provider_account_plan_type: Some("pro".to_owned()),
provider_account_plan_type_display: Some("Pro".to_owned()),
provider_account_authentication_kind: Some("api_key".to_owned()),
upstream_model_id: Some("upstream-err".to_owned()),
upstream_transport: Some("http_sse".to_owned()),
failure_kind: "upstream_error".to_owned(),
upstream_send_state: Some("sent".to_owned()),
client_status_code: Some(502),
upstream_status_code: Some(502),
provider_error_code: Some("upstream".to_owned()),
client_response_id: None,
upstream_request_id: None,
latency_ms: Some(90),
message: "snapshot error".to_owned(),
error_details: Some(
r#"{"error":{"code":"upstream","message":"raw upstream marker"}}"#.to_owned(),
),
client_ip: Some("203.0.113.8".to_owned()),
user_agent: Some("codex-cli/0.144.0".to_owned()),
reasoning_effort: Some("medium".to_owned()),
reasoning_preset: None,
request_kind: Some("root".to_owned()),
subagent_kind: None,
compact: Some(false),
continuation_affinity_hash: Some("affinity-hash".to_owned()),
continuation_previous_response_id_hash: Some("response-hash".to_owned()),
continuation_unavailable_reason: Some("reused_connection_lost".to_owned()),
upstream_connection_id: Some("connection-id".to_owned()),
upstream_connection_exit_reason: Some("tcp_reset".to_owned()),
upstream_connection_age_ms: Some(12_000),
upstream_connection_idle_ms: Some(4_000),
recovery_request_id: None,
recovered_at: None,
recovery_attempt_count: 0,
recovery_retry_delay_ms: None,
recovery_total_latency_ms: None,
occurred_at: Utc::now(),
stable_sort_id: "model_request:req_err".to_owned(),
});
.push(ops_error_record());
let response = observability::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
@@ -621,7 +564,7 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
"userAgent": "codex-cli/0.144.0",
"providerErrorCode": "upstream",
"upstreamSendState": "sent",
"errorDetails": "{\"error\":{\"code\":\"upstream\",\"message\":\"raw upstream marker\"}}",
"errorDetails": serde_json::to_string_pretty(&json!({"error": {"code": "upstream", "message": "raw upstream marker"}})).unwrap(),
"continuationUnavailableReason": "reused_connection_lost",
"upstreamConnectionExitReason": "tcp_reset",
"upstreamConnectionAgeMs": 12000,
@@ -707,7 +650,7 @@ async fn diagnostics_should_keep_stable_key_and_display_name_contract() {
);
}
fn usage_record_with_account(
pub(super) fn usage_record_with_account(
id: &str,
account_id: &str,
name: &str,
@@ -858,6 +801,7 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
.lock()
.expect("usage records")
.push(UsageListRecord {
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
client_api_key_name: Some("Production".to_owned()),
id: "request_endpoint".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -1303,3 +1247,123 @@ async fn observability_queries_reject_group_filters_and_analytics_search() {
assert!(observations.records.is_empty());
assert!(observations.errors.is_empty());
}
pub(super) fn ops_error_record() -> gateway_admin::model::observability::OpsError {
use gateway_admin::model::observability::OpsError;
OpsError {
client_api_key_name: Some("Production".to_owned()),
source: "model_request".to_owned(),
event_id: "err_snapshot".to_owned(),
request_id: Some("req_err".to_owned()),
attempt_index: Some(1),
client_api_key_ref: Some("key_err".to_owned()),
component: "model_request".to_owned(),
operation: "responses".to_owned(),
protocol: Some("openai".to_owned()),
client_transport: Some("http_sse".to_owned()),
requested_model_id: Some("gpt-5.4".to_owned()),
service_tier: Some("priority".to_owned()),
endpoint: Some("/v1/responses".to_owned()),
provider_kind: Some("openai".to_owned()),
provider_account_ref: Some("acct_err".to_owned()),
provider_account_name: None,
provider_account_email: Some("err@example.invalid".to_owned()),
provider_account_plan_type: Some("pro".to_owned()),
provider_account_plan_type_display: Some("Pro".to_owned()),
provider_account_authentication_kind: Some("api_key".to_owned()),
upstream_model_id: Some("upstream-err".to_owned()),
upstream_transport: Some("http_sse".to_owned()),
failure_kind: "upstream_error".to_owned(),
upstream_send_state: Some("sent".to_owned()),
client_status_code: Some(502),
upstream_status_code: Some(502),
provider_error_code: Some("upstream".to_owned()),
client_response_id: None,
upstream_request_id: None,
latency_ms: Some(90),
message: "snapshot error".to_owned(),
error_details: Some(
r#"{"error":{"code":"upstream","message":"raw upstream marker"}}"#.to_owned(),
),
client_ip: Some("203.0.113.8".to_owned()),
user_agent: Some("codex-cli/0.144.0".to_owned()),
reasoning_effort: Some("medium".to_owned()),
reasoning_preset: None,
request_kind: Some("root".to_owned()),
subagent_kind: None,
compact: Some(false),
continuation_affinity_hash: Some("affinity-hash".to_owned()),
continuation_previous_response_id_hash: Some("response-hash".to_owned()),
continuation_unavailable_reason: Some("reused_connection_lost".to_owned()),
upstream_connection_id: Some("connection-id".to_owned()),
upstream_connection_exit_reason: Some("tcp_reset".to_owned()),
upstream_connection_age_ms: Some(12_000),
upstream_connection_idle_ms: Some(4_000),
recovery_request_id: None,
recovered_at: None,
recovery_attempt_count: 0,
recovery_retry_delay_ms: None,
recovery_total_latency_ms: None,
occurred_at: Utc::now(),
stable_sort_id: "model_request:req_err".to_owned(),
}
}
#[tokio::test]
async fn usage_status_query_accepts_three_modes_and_rejects_unknown_values() {
use crate::admin::{AdminTestFixture, AdminTestState};
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header},
};
use gateway_api::admin::observability;
use tower::ServiceExt as _;
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
for (status, expected) in [
("success", StatusCode::OK),
("error", StatusCode::OK),
("all", StatusCode::OK),
("failed", StatusCode::BAD_REQUEST),
] {
let response = observability::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
Request::builder()
.uri(format!("/api/admin/usage/records?status={status}"))
.header(header::COOKIE, "cpr_session=valid-session")
.header("x-request-id", "usage_status_fixture")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let actual = response.status();
let body = to_bytes(response.into_body(), 1024 * 1024).await.unwrap();
assert_eq!(
actual,
expected,
"{status}: {}",
String::from_utf8_lossy(&body)
);
if expected == StatusCode::OK {
use gateway_admin::model::observability::UsageRecordStatus;
let status_value = match status {
"success" => UsageRecordStatus::Success,
"error" => UsageRecordStatus::Error,
_ => UsageRecordStatus::All,
};
assert_eq!(
fixture
.observations
.lock()
.unwrap()
.records
.last()
.unwrap()
.status,
status_value
);
}
}
}
@@ -1,5 +1,7 @@
//! 插件管理 HTTP 测试入口,以及制品、实例与权限合同测试
mod usage_data;
use async_trait::async_trait;
use gateway_admin::model::plugins::distribution::{
DownloadedPlugin, GithubReleaseQuery, PluginRelease, RemotePluginLocation, SourceCredential,
@@ -0,0 +1,24 @@
//! 通用用量端口的插件管理测试夹具,非本测试调用路径禁止访问
use super::TestPluginPorts;
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
#[async_trait::async_trait]
impl UsageDataStore for TestPluginPorts {
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
unreachable!("unexpected usage call")
}
async fn list_records(&self, _: ListUsageRecords) -> UsageDataResult<UsageRecordPage> {
unreachable!("unexpected usage call")
}
async fn write_records(&self, _: StoredUsageWrite) -> UsageDataResult<UsageWriteResult> {
unreachable!("unexpected usage call")
}
async fn list_aggregates(&self, _: ListUsageAggregates) -> UsageDataResult<UsageAggregatePage> {
unreachable!("unexpected usage call")
}
async fn write_aggregates(
&self,
_: StoredAggregateWrite,
) -> UsageDataResult<UsageAggregateWriteResult> {
unreachable!("unexpected usage call")
}
}
@@ -62,7 +62,7 @@ fn update_body() -> Value {
"smartScheduling": gateway_core::account::SmartSchedulingConfig::default(),
"minCodexDesktopVersion": "26.825.6671",
"minCodexCliVersion": "0.40.0",
"usageRetentionDays": 32,
"usageRetentionDays": 31,
"opsEventRetentionDays": 31,
"auditRetentionDays": 91,
"accountAutoFreezeEnabled": true,
@@ -141,7 +141,7 @@ async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_int
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3, "preferHigherWeight": true});
let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3});
body["smartScheduling"] = custom.clone();
let response = app(fixture.state())
.oneshot(request(
@@ -161,8 +161,8 @@ async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_int
let mut invalid_values = vec![
json!(null),
json!({}),
json!({"loadWeight":0,"quotaWeight":0,"healthWeight":0,"latencyWeight":0,"resetWeight":0,"queueWeight":0,"preferHigherWeight":true}),
json!({"loadWeight":0.01,"quotaWeight":1,"healthWeight":1,"latencyWeight":1,"resetWeight":0,"queueWeight":0,"preferHigherWeight":false}),
json!({"loadWeight":0,"quotaWeight":0,"healthWeight":0,"latencyWeight":0,"resetWeight":0,"queueWeight":0}),
json!({"loadWeight":0.01,"quotaWeight":1,"healthWeight":1,"latencyWeight":1,"resetWeight":0,"queueWeight":0}),
];
for field in ["resetWeight", "queueWeight"] {
for value in [json!(-0.1), json!(10.1), json!(0.01), json!(null)] {
@@ -216,6 +216,56 @@ fn settings_request_should_reject_unknown_rotation_strategy() {
assert_eq!(request.validate().unwrap_err().field(), "rotationStrategy");
}
#[tokio::test]
async fn short_request_retention_round_trips_and_invalid_days_do_not_replace_it() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
for days in [1, 3, 7, 15, 30, 31] {
body["usageRetentionDays"] = json!(days);
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(body.clone()),
))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let saved = response_json(response).await;
assert_eq!(saved["data"]["usageRetentionDays"], days);
body["configRevision"] = saved["data"]["configRevision"].clone();
}
for (days, expected_status) in [
(json!(0), StatusCode::BAD_REQUEST),
(json!(32), StatusCode::BAD_REQUEST),
(json!(u32::MAX), StatusCode::BAD_REQUEST),
(json!(-1), StatusCode::UNPROCESSABLE_ENTITY),
(json!(1.5), StatusCode::UNPROCESSABLE_ENTITY),
(json!(null), StatusCode::UNPROCESSABLE_ENTITY),
] {
body["usageRetentionDays"] = days;
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(body.clone()),
))
.await
.unwrap();
assert_eq!(response.status(), expected_status);
}
let response = app(fixture.state())
.oneshot(request(Method::GET, "/api/admin/settings", None))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response_json(response).await["data"]["usageRetentionDays"],
31
);
}
#[test]
fn settings_request_accepts_unlimited_default_account_concurrency() {
let mut body = update_body();
@@ -298,7 +348,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
min_codex_desktop_version: Some("26.825.6671".to_owned()),
min_codex_cli_version: Some("0.40.0".to_owned()),
usage_retention_days: 32,
usage_retention_days: 31,
ops_event_retention_days: 31,
audit_retention_days: 91,
account_auto_freeze_enabled: true,
@@ -349,7 +399,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
"smartSchedulingDefaults": gateway_core::account::SmartSchedulingConfig::default(),
"minCodexDesktopVersion": "26.825.6671",
"minCodexCliVersion": "0.40.0",
"usageRetentionDays": 32,
"usageRetentionDays": 31,
"opsEventRetentionDays": 31,
"auditRetentionDays": 91,
"accountAutoFreezeEnabled": true,
@@ -93,6 +93,7 @@ pub(super) async fn fixture() -> AdminTestFixture {
fn usage_record() -> UsageListRecord {
UsageListRecord {
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
client_api_key_name: Some("Production".to_owned()),
id: "req-visible".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -0,0 +1,112 @@
//! 真实 HTTP / WebSocket 路由与 Core 执行入口共同拒绝生图,不触发准入或上游
use super::*;
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header::AUTHORIZATION},
};
use futures::{SinkExt, StreamExt};
use serde_json::{Value, json};
use std::{sync::Arc, time::Duration};
use tokio_tungstenite::tungstenite::{Message, client::IntoClientRequest};
use tower::ServiceExt;
fn execution() -> Arc<DefaultExecutionService> {
Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot_with_image_policy(
"sk_image_policy",
"openai",
"key_image_policy",
&["acct_api_test"],
false,
)),
Arc::new(UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
Arc::new(crate::support::RecordingDiagnostics::default()),
))
}
#[tokio::test]
async fn denied_image_requests_return_403_before_admission_or_upstream() {
let app = api_router(execution()).await;
for (path, body) in [
(
"/v1/images/generations",
json!({"model":"gpt-image-2","prompt":"test"}),
),
(
"/v1/images/edits",
json!({"model":"gpt-image-2","prompt":"test","images":[{"image_url":"data:image/png;base64,AA=="}]}),
),
(
"/v1/responses",
json!({"model":"model-a","input":"test","tools":[{"type":"image_generation"}]}),
),
(
"/v1/responses",
json!({"model":"model-a","input":"test","stream":true,"tool_choice":{"type":"image_generation"}}),
),
] {
let response = app
.clone()
.oneshot(
Request::post(path)
.header(AUTHORIZATION, "Bearer sk_image_policy")
.header("content-type", "application/json")
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body: Value =
serde_json::from_slice(&to_bytes(response.into_body(), usize::MAX).await.unwrap())
.unwrap();
assert_eq!(status, StatusCode::FORBIDDEN, "{path}: {body}");
assert_eq!(body["error"]["code"], "image_generation_disabled");
}
}
#[tokio::test]
async fn each_websocket_create_enforces_image_policy() {
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 = 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, "Bearer sk_image_policy".parse().unwrap());
let (mut socket, _) = tokio_tungstenite::connect_async(request).await.unwrap();
for image_fields in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
] {
let mut body = json!({"type":"response.create","model":"model-a","input":"test"});
body.as_object_mut()
.unwrap()
.extend(image_fields.as_object().unwrap().clone());
socket
.send(Message::Text(body.to_string().into()))
.await
.unwrap();
let frame = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
let body: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(body["type"], "error", "{body}");
assert_eq!(body["status"], 403);
assert_eq!(body["error"]["code"], "image_generation_disabled");
}
socket.close(None).await.unwrap();
server.abort();
}
+18 -4
View File
@@ -3,6 +3,7 @@
mod auth;
mod endpoint;
mod error;
mod image_policy;
mod images;
mod live;
mod middleware;
@@ -292,6 +293,16 @@ fn snapshot_with_accounts(
provider_name: &str,
client_key_id: &str,
accounts: &[&str],
) -> RuntimeSnapshot {
snapshot_with_image_policy(plaintext, provider_name, client_key_id, accounts, true)
}
fn snapshot_with_image_policy(
plaintext: &str,
provider_name: &str,
client_key_id: &str,
accounts: &[&str],
allow_image_generation: bool,
) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from_iter(
@@ -323,10 +334,13 @@ fn snapshot_with_accounts(
vec![ClientPolicy::new(
ClientApiKeyId::new(client_key_id).expect("key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext key"),
Arc::new(FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)),
Arc::new(
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)
.with_image_generation(allow_image_generation),
),
true,
RateLimits::unlimited(),
)],
@@ -247,6 +247,7 @@ impl ClientRoutingScope {
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
allow_image_generation: bool,
fast_mode: FastMode,
request_profiles: BTreeMap<ProviderKind, super::OpaqueProviderData>,
directory: Arc<RuntimeAccountDirectory>,
@@ -257,6 +258,18 @@ pub struct FrozenAccountScope {
}
impl FrozenAccountScope {
/// Key 绑定分组合并后的生图权限,与账号成员资格一同冻结
#[must_use]
pub const fn with_image_generation(mut self, allowed: bool) -> Self {
self.allow_image_generation = allowed;
self
}
#[must_use]
pub const fn allows_image_generation(&self) -> bool {
self.allow_image_generation
}
/// 与 Key 授权范围一同冻结;具体字段只由对应 Provider 解释
#[must_use]
pub fn with_request_profiles(
@@ -296,6 +309,7 @@ impl FrozenAccountScope {
ClientRoutingScope::Restricted { provider_kinds, .. } => Arc::clone(provider_kinds),
};
Self {
allow_image_generation: true,
fast_mode: FastMode::Default,
request_profiles: BTreeMap::new(),
directory,
@@ -510,6 +510,8 @@ pub struct AccountSelectionContext {
pub policy: AccountSelectionPolicy,
pub now: SystemTime,
pub excluded_accounts: BTreeSet<ProviderAccountId>,
/// 租约争用或队列让位只阻止立即选号,不排除等待资格
pub busy_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
/// 优先账号可调度时,是否优先于其它账号的持久权重
pub preferred_account_overrides_weight: bool,
@@ -602,7 +604,7 @@ impl<'a> AccountSelection<'a> {
pub struct AccountSelector;
impl AccountSelector {
/// 返回调度策略可见的全部合格候选;权重层授权由调用策略的宿主适配器裁剪
/// 策略同时约束立即选号与等待,候选保留仅因本地容量受阻的账号
#[must_use]
pub(crate) fn policy_candidates<'a>(
&self,
@@ -611,7 +613,12 @@ impl AccountSelector {
) -> Vec<&'a AccountCandidate> {
candidates
.iter()
.filter(|candidate| self.scheduling_blocker(candidate, context).is_none())
.filter(|candidate| {
matches!(
self.scheduling_blocker(candidate, context),
None | Some(AccountSchedulingBlocker::ConcurrencyLimit)
)
})
.collect()
}
@@ -697,10 +704,7 @@ impl AccountSelector {
.max()?;
let (mut preferred, preferred_candidate) = self.preferred_decision(candidates, context);
if let Some(candidate) = preferred_candidate {
// 权重回切只裁决内置策略的软亲和,插件显式选号按实际选择记录结果
let prefer_higher_weight = context.policy.strategy() == RotationStrategy::Smart
&& context.policy.smart_scheduling().prefer_higher_weight();
if (!context.preferred_account_overrides_weight || prefer_higher_weight)
if !context.preferred_account_overrides_weight
&& candidate.account.weight() < highest_weight
{
preferred =
@@ -860,6 +864,9 @@ impl AccountSelector {
if context.excluded_accounts.contains(candidate.account.id()) {
return Some(AccountSchedulingBlocker::Excluded);
}
if context.busy_accounts.contains(candidate.account.id()) {
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
}
if context
.concurrency_limit(&candidate.account)
.limit()
@@ -6,7 +6,6 @@ use serde::{Deserialize, Serialize};
#[serde(into = "SmartSchedulingValues")]
pub struct SmartSchedulingConfig {
weights: [u8; 6],
prefer_higher_weight: bool,
}
impl<'de> Deserialize<'de> for SmartSchedulingConfig {
@@ -24,7 +23,6 @@ impl Default for SmartSchedulingConfig {
fn default() -> Self {
Self {
weights: [10, 8, 10, 5, 0, 0],
prefer_higher_weight: false,
}
}
}
@@ -39,10 +37,7 @@ pub enum SmartSchedulingConfigError {
impl SmartSchedulingConfig {
/// 顺序为负载、剩余额度、健康、首输出延迟、额度重置和排队压力
pub fn new(
weights: [f64; 6],
prefer_higher_weight: bool,
) -> Result<Self, SmartSchedulingConfigError> {
pub fn new(weights: [f64; 6]) -> Result<Self, SmartSchedulingConfigError> {
let mut tenths = [0; 6];
for (index, field) in [
"loadWeight",
@@ -67,10 +62,7 @@ impl SmartSchedulingConfig {
if tenths == [0; 6] {
return Err(SmartSchedulingConfigError::EmptyWeights);
}
Ok(Self {
weights: tenths,
prefer_higher_weight,
})
Ok(Self { weights: tenths })
}
#[must_use]
@@ -78,11 +70,6 @@ impl SmartSchedulingConfig {
self.weights.map(|weight| f64::from(weight) / 10.0)
}
#[must_use]
pub const fn prefer_higher_weight(self) -> bool {
self.prefer_higher_weight
}
pub(crate) fn score_tolerance(self) -> f64 {
// 与系数同步缩放,避免仅放大相同比例的系数就改变近似最优候选集合
// 排队系数只参与入队选择,不能放大立即选号的候选容差
@@ -105,24 +92,20 @@ struct SmartSchedulingValues {
latency_weight: f64,
reset_weight: f64,
queue_weight: f64,
prefer_higher_weight: bool,
}
impl TryFrom<SmartSchedulingValues> for SmartSchedulingConfig {
type Error = SmartSchedulingConfigError;
fn try_from(value: SmartSchedulingValues) -> Result<Self, Self::Error> {
Self::new(
[
value.load_weight,
value.quota_weight,
value.health_weight,
value.latency_weight,
value.reset_weight,
value.queue_weight,
],
value.prefer_higher_weight,
)
Self::new([
value.load_weight,
value.quota_weight,
value.health_weight,
value.latency_weight,
value.reset_weight,
value.queue_weight,
])
}
}
@@ -143,7 +126,6 @@ impl From<SmartSchedulingConfig> for SmartSchedulingValues {
latency_weight,
reset_weight,
queue_weight,
prefer_higher_weight: value.prefer_higher_weight,
}
}
}
@@ -132,6 +132,15 @@ impl TraceContext {
}
let data = json!({
"kind": error.kind().as_str(), "sendState": error.send_state().as_str(),
"recovery": {
"replaySafe": error.replay_is_safe(),
"retryProhibited": error.retry_is_prohibited(),
"sameAccountCredentialRetry": error.retries_same_account(),
"intent": error.pre_delivery_retry().map(|intent| format!("{intent:?}")),
"continuationFailure": error.continuation_failure().map(|failure| format!("{failure:?}")),
"continuationDisposition": error.continuation_recovery_disposition().map(|disposition| format!("{disposition:?}")),
"retryAfterMs": error.retry_after().map(|delay| delay.as_millis()),
},
"diagnostic": error.diagnostic().map(|diagnostic| json!({
"stage": diagnostic.stage(), "code": diagnostic.code(), "message": diagnostic.as_str(),
})),
@@ -291,6 +300,27 @@ impl TraceContext {
}
summary["upstreamStatus"] =
json!(data.get("upstreamStatus").and_then(Value::as_u64));
if let Some(recovery) = data.get("recovery") {
let mut kept = json!({});
for key in [
"replaySafe",
"retryProhibited",
"sameAccountCredentialRetry",
] {
kept[key] = json!(recovery.get(key).and_then(Value::as_bool));
}
for key in ["intent", "continuationFailure", "continuationDisposition"] {
kept[key] = json!(
recovery
.get(key)
.and_then(Value::as_str)
.map(|value| bounded(value, 256))
);
}
kept["retryAfterMs"] =
json!(recovery.get("retryAfterMs").and_then(Value::as_u64));
summary["recovery"] = kept;
}
if let Some(diagnostic) = data.get("diagnostic").and_then(Value::as_object) {
let mut kept = json!({});
for (key, limit) in [("stage", 96), ("code", 96), ("message", 256)] {
@@ -359,6 +389,7 @@ impl TraceContext {
| "account.selected"
| "attempt.failed"
| "retry.decided"
| "retry.policy"
| "upstream.close"
| "upstream.read.failed"
| "capture.gap"
@@ -294,6 +294,20 @@ where
"admissionMs": request.admission_decision_ms,
"continuation": format!("{continuation_attempt:?}"),
"candidateCount": plan.candidates().len(),
"accountPolicy": {
"strategy": plan.account_selection_policy().strategy().as_str(),
"affinity": plan.account_selection_policy().openai_account_affinity().as_str(),
"maxAccountRotations": plan.account_selection_policy().max_account_rotations(),
"maxRoutingAttempts": plan.max_attempts().get(),
"maxConcurrentPerAccount": plan.account_selection_policy().max_concurrent_per_account().get(),
"queueMaxWaiting": plan.account_selection_policy().queue_policy().max_waiting,
"queueTimeoutMs": plan.account_selection_policy().queue_policy().timeout.as_millis(),
"sessionAffinityTtlMs": plan.account_selection_policy().openai_session_affinity_ttl().as_millis(),
"smartWeights": plan.account_selection_policy().smart_scheduling().weights(),
},
"requiredAccount": account_selection.required_account().map(|account| account.as_str()),
"hasMiddleware": middleware.is_some(),
"hasRetryPolicy": request_policy.is_some(),
}),
);
let image_generation_requested = operation.image_generation_requested();
@@ -1200,6 +1214,17 @@ where
retry_allowed: continuation_retry || candidate_retry,
})
.await?;
attempt_trace.record("retry.decided", json!({
"phase": "prepare", "retryable": retryable,
"continuationRetry": continuation_retry, "candidateRetry": candidate_retry,
"retryProhibited": error.retry_is_prohibited(),
"sendState": error.send_state().as_str(),
"requestSendState": self.current_send_state().as_str(),
"continuation": format!("{:?}", self.continuation_attempt),
"accountRotations": self.account_rotations,
"maxAccountRotations": self.plan.account_selection_policy().max_account_rotations(),
"remainingRoutingAttempts": self.plan.max_attempts().get().saturating_sub(self.routing_attempts),
}));
if retryable {
return Ok(Some(PullOutcome::AttemptDiscarded));
}
@@ -1706,6 +1731,7 @@ where
};
self.trace.attempt(current.index.get()).record("retry.decided", json!({
"phase": "execute",
"retryable": retryable, "continuationRetry": continuation_retry,
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
@@ -1719,6 +1745,17 @@ where
"delayMs": transient_retry.or(transport_recovery.map(|(_, delay)| delay)).map(duration_ms),
"downstreamCommitted": self.response_commit.committed_at().is_some(),
"replayBarrierReached": self.replay_barrier_reached,
"deliveryPending": self.delivery_pending,
"providerProvedReplaySafe": provider_proved_replay_safe,
"retryProhibited": error.retry_is_prohibited(),
"requiredAccount": self.account_selection.required_account().map(|account| account.as_str()),
"continuation": format!("{:?}", self.continuation_attempt),
"connectionBudgetExhausted": self.connection_budget.exhausted(),
"remainingRoutingAttempts": self.plan.max_attempts().get().saturating_sub(self.routing_attempts),
"maxAccountRotations": self.plan.account_selection_policy().max_account_rotations(),
"transientRetriesUsed": self.transient_retry_counts.get(current.metadata.provider_account_id()).copied().unwrap_or_default(),
"credentialRecoveryAttempted": self.credential_recovery_attempted_accounts.contains(current.metadata.provider_account_id()),
"hasRetryPolicy": self.request_policy.is_some(),
"sendState": format!("{attempt_send_state:?}"),
}));
if retryable {
@@ -1818,6 +1855,7 @@ where
&& !self.cancellation.is_cancelled()
&& !self.execution_effect_observed();
let allowed = facts.retry_allowed;
let attempt_index = facts.attempt_index;
let decision = {
let cancellation = self.cancellation.clone();
let cancelled = cancellation.cancelled().fuse();
@@ -1830,6 +1868,16 @@ where
decision = decision => Ok(decision),
}
}?;
self.trace.attempt(attempt_index.get()).record(
"retry.policy",
json!({
"hostAllowed": allowed,
"policyStopped": decision == super::policy::RetryDecision::Stop,
"executionEffectObserved": self.execution_effect_observed(),
"cancelled": self.cancellation.is_cancelled(),
"deadlineElapsed": self.deadline.is_elapsed(),
}),
);
// 策略只收窄宿主已允许的恢复,返回后仍复核期限和外部副作用
Ok(allowed
&& decision != super::policy::RetryDecision::Stop
@@ -400,6 +400,10 @@ impl DefaultExecutionService {
started_at: SystemTime,
request_id: Option<ModelRequestId>,
) -> Result<StartedExecution, GatewayError> {
crate::engine::image_generation::check_image_generation(
authorization.account_scope.allows_image_generation(),
&request.operation,
)?;
let request_id = request_id.map_or_else(new_request_id, Ok)?;
if authorization.cancellation.is_cancelled() {
return Err(GatewayError::new(
@@ -0,0 +1,20 @@
//! 冻结分组生图权限的请求与发送前门禁
use crate::{
error::{GatewayError, GatewayErrorKind},
operation::Operation,
};
pub(super) fn check_image_generation(
allowed: bool,
operation: &Operation,
) -> Result<(), GatewayError> {
if !allowed && operation.image_generation_requested() {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"Image generation is not allowed by the API key's account groups",
)
.with_client_code("image_generation_disabled"));
}
Ok(())
}
+25 -2
View File
@@ -9,6 +9,7 @@ pub mod continuation;
pub mod coordinator;
pub mod execution;
pub mod extensions;
mod image_generation;
pub mod middleware;
pub mod nested;
pub mod observation;
@@ -37,7 +38,9 @@ use crate::account::{
ProviderAccountId,
};
use crate::engine::continuation::{ContinuationBinding, NativeContinuationPin};
use crate::error::{GatewayError, ProviderConnectionObservation, ProviderError, StoreError};
use crate::error::{
GatewayError, ProviderConnectionObservation, ProviderError, ProviderErrorKind, StoreError,
};
use crate::event::ProviderEvent;
use crate::identity::ProviderKind;
use crate::lifecycle::CancellationToken;
@@ -688,6 +691,26 @@ impl AttemptContext {
account_id: ProviderAccountId,
terminal: provider::ProviderMiddlewareTerminal,
) -> Result<provider::ProviderStream, ProviderError> {
// 入口已校验原请求,此处复查中间件最终提交的 operation,每次重试均生效
let terminal: provider::ProviderMiddlewareTerminal = if self
.account_scope()
.is_some_and(|scope| !scope.allows_image_generation())
{
Box::new(move |operation, headers| {
Box::pin(async move {
image_generation::check_image_generation(false, &operation).map_err(|_| {
ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::NotSent,
)
.with_retry_prohibited()
})?;
terminal(operation, headers).await
})
})
} else {
terminal
};
let context = middleware::MiddlewareContext::new(
middleware::MiddlewareTarget {
request_id: self.request.request_id.clone(),
@@ -806,7 +829,7 @@ impl AttemptContext {
provider: &ProviderKind,
model: Option<&str>,
candidates: &'a [AccountCandidate],
context: &AccountSelectionContext,
context: &mut AccountSelectionContext,
) -> Result<Option<AccountSelection<'a>>, policy::AccountPolicyError> {
match self.request.request_policy() {
Some(policy) => {
@@ -113,6 +113,7 @@ pub enum ModelRouteDecision {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountScheduleCandidate {
account_id: ProviderAccountId,
available: bool,
weight: u16,
in_flight: u32,
maximum_concurrency: u32,
@@ -124,6 +125,11 @@ pub struct AccountScheduleCandidate {
}
impl AccountScheduleCandidate {
#[must_use]
pub const fn available(&self) -> bool {
self.available
}
#[must_use]
pub const fn account_id(&self) -> &ProviderAccountId {
&self.account_id
@@ -242,6 +248,7 @@ impl AccountScheduleInput {
pub enum AccountScheduleDecision {
Pick(ProviderAccountId),
Delegate,
DelegateFiltered(BTreeSet<ProviderAccountId>),
Reject,
}
@@ -284,6 +291,8 @@ pub struct RequestPolicyFault;
/// Provider 选择器需要区分明确拒绝、策略故障和等待期间过期的选择
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum AccountPolicyError {
#[error("account scheduling policy left no selectable account")]
Filtered,
#[error("account scheduling policy rejected the request")]
Rejected,
#[error("account scheduling policy failed")]
@@ -414,14 +423,14 @@ impl RequestPolicyContext {
.await
}
/// 硬绑定已在候选资格中收窄;插件先决定,委托时再采用内置亲和与排序
/// 插件只处理宿主候选;返回后的上下文保留过滤结果,供亲和重选和等待共用
pub async fn select_account<'a>(
&self,
attempt_index: NonZeroU32,
provider: &ProviderKind,
model: Option<&str>,
candidates: &'a [AccountCandidate],
context: &AccountSelectionContext,
context: &mut AccountSelectionContext,
) -> Result<Option<AccountSelection<'a>>, AccountPolicyError> {
let eligible = AccountSelector.policy_candidates(candidates, context);
if eligible.is_empty() {
@@ -431,6 +440,9 @@ impl RequestPolicyContext {
.iter()
.map(|candidate| AccountScheduleCandidate {
account_id: candidate.account.id().clone(),
available: AccountSelector
.scheduling_blocker(candidate, context)
.is_none(),
weight: candidate.account.weight().get(),
in_flight: candidate.signals.in_flight,
maximum_concurrency: context.concurrency_limit(&candidate.account).get(),
@@ -458,27 +470,47 @@ impl RequestPolicyContext {
.map_err(|_| AccountPolicyError::Fault)?;
// 插件调用可能消耗显著时间;无论插件选择还是委托,均用返回时刻重新执行
// 同一资格判断,不能让调用前尚未过期的候选在调用后被内置策略选中
let mut current = context.clone();
current.now = SystemTime::now();
match decision {
AccountScheduleDecision::Delegate => AccountSelector
.select(candidates, &current)
.map(Some)
.ok_or(AccountPolicyError::StaleCandidate),
AccountScheduleDecision::Reject => Err(AccountPolicyError::Rejected),
AccountScheduleDecision::Pick(account_id) => {
if !eligible
AccountScheduleDecision::Delegate => {}
AccountScheduleDecision::Reject => return Err(AccountPolicyError::Rejected),
AccountScheduleDecision::DelegateFiltered(excluded) => {
let candidate_ids = eligible
.iter()
.any(|candidate| candidate.account.id() == &account_id)
.map(|candidate| candidate.account.id())
.collect::<std::collections::HashSet<_>>();
if excluded.iter().any(|id| !candidate_ids.contains(id)) {
return Err(AccountPolicyError::Fault);
}
let all_filtered = eligible
.iter()
.all(|candidate| excluded.contains(candidate.account.id()));
context.excluded_accounts.extend(excluded);
if all_filtered {
return Err(AccountPolicyError::Filtered);
}
}
AccountScheduleDecision::Pick(account_id) => {
if AccountSelector
.select_policy_candidate(candidates, context, &account_id)
.is_none()
{
return Err(AccountPolicyError::Fault);
}
AccountSelector
.select_policy_candidate(candidates, &current, &account_id)
context.now = SystemTime::now();
return AccountSelector
.select_policy_candidate(candidates, context, &account_id)
.map(Some)
.ok_or(AccountPolicyError::StaleCandidate)
.ok_or(AccountPolicyError::StaleCandidate);
}
}
context.now = SystemTime::now();
if AccountSelector
.policy_candidates(candidates, context)
.is_empty()
{
return Err(AccountPolicyError::StaleCandidate);
}
Ok(AccountSelector.select(candidates, context))
}
}
+37 -8
View File
@@ -403,6 +403,7 @@ pub struct GenerateRequest {
#[derive(Clone, PartialEq)]
struct GeneratePayload {
source_image_generation: bool,
protocol_payload: ProtocolPayload,
client_turn: Option<ClientTurn>,
provider_session_state: Option<ProviderSessionState>,
@@ -445,6 +446,7 @@ impl GenerateRequest {
pub fn from_protocol_payload(protocol_payload: ProtocolPayload) -> Self {
Self {
payload: Arc::new(GeneratePayload {
source_image_generation: false,
protocol_payload,
client_turn: None,
provider_session_state: None,
@@ -506,14 +508,9 @@ impl GenerateRequest {
/// 返回客户端是否请求了图片生成工具
#[must_use]
pub fn image_generation_requested(&self) -> bool {
self.body()
.get("tools")
.and_then(Value::as_array)
.is_some_and(|tools| {
tools.iter().any(|tool| {
tool.get("type").and_then(Value::as_str) == Some("image_generation")
})
})
self.payload.source_image_generation
|| (self.protocol_payload().protocol() == "openai"
&& requests_image_generation(self.body()))
}
/// 返回指定 Provider 的连接内会话状态
@@ -612,6 +609,35 @@ fn contains_type(value: Option<&Value>, expected: &str) -> bool {
}
}
/// 原生工具声明授权上游生成图片,Codex 的被动 namespace 与图片输入不属于该权限
fn requests_image_generation(body: &Map<String, Value>) -> bool {
contains_image_tool(body.get("tools"))
|| body
.get("input")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type").and_then(Value::as_str) == Some("additional_tools")
&& contains_image_tool(item.get("tools"))
})
})
|| body.get("tool_choice").is_some_and(|choice| {
is_image_tool(choice)
|| (choice.get("type").and_then(Value::as_str) == Some("allowed_tools")
&& contains_image_tool(choice.get("tools")))
})
}
fn contains_image_tool(value: Option<&Value>) -> bool {
value
.and_then(Value::as_array)
.is_some_and(|tools| tools.iter().any(is_image_tool))
}
fn is_image_tool(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("image_generation")
}
/// 图像 API 的稳定端点语义
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ImageRequestKind {
@@ -1020,6 +1046,7 @@ impl Operation {
Self::Generate(request) => {
let protocol_payload = request.protocol_payload();
let source_requirements = translating.then(|| request.requirements());
let source_image_generation = translating && request.image_generation_requested();
let body = match body {
Some(body) => {
serde_json::from_slice::<Map<String, Value>>(&body).map_err(|_| {
@@ -1038,6 +1065,8 @@ impl Operation {
};
let mut payload = (*request.payload).clone();
payload.protocol_payload = replacement;
// 跨协议转换不能抹去已识别的生图权限要求,也不解释目标私有协议
payload.source_image_generation |= source_image_generation;
if let Some(requirements) = source_requirements {
payload.source_requirements = Some(requirements);
}
@@ -74,6 +74,7 @@ impl SnapshotClientPolicyFacts {
/// Store 读取到的账号分组事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupFacts {
allow_image_generation: bool,
fast_mode: FastMode,
id: AccountGroupId,
name: String,
@@ -81,6 +82,12 @@ pub struct SnapshotAccountGroupFacts {
}
impl SnapshotAccountGroupFacts {
#[must_use]
pub const fn with_image_generation(mut self, allowed: bool) -> Self {
self.allow_image_generation = allowed;
self
}
#[must_use]
pub const fn with_fast_mode(mut self, fast_mode: FastMode) -> Self {
self.fast_mode = fast_mode;
@@ -93,6 +100,7 @@ impl SnapshotAccountGroupFacts {
id,
name,
enabled,
allow_image_generation: true,
fast_mode: FastMode::Default,
}
}
@@ -499,6 +507,7 @@ async fn compile_runtime_snapshot(
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
for policy in facts.client_policies {
let mut fast_mode = FastMode::Default;
let mut allow_image_generation = true;
let account_scope = if policy.group_ids.is_empty() {
FrozenAccountScope::new(
Arc::clone(&account_directory),
@@ -515,8 +524,9 @@ async fn compile_runtime_snapshot(
let group = groups
.get(&group_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
// 禁用分组仅影响选号;Key 仍绑定其 Fast 策略
// 停用只影响选号,仍绑定的分组继续贡献请求策略
fast_mode = fast_mode.merge(group.fast_mode);
allow_image_generation &= group.allow_image_generation;
bound_groups.push(RoutingGroupSnapshot::new(
group.id.clone(),
group.name.clone(),
@@ -539,6 +549,7 @@ async fn compile_runtime_snapshot(
policy.plaintext_key,
Arc::new(
account_scope
.with_image_generation(allow_image_generation)
.with_fast_mode(fast_mode)
.with_request_profiles(policy.request_profiles),
),
@@ -102,6 +102,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
policy: AccountSelectionPolicy::new(strategy, NonZeroU32::new(3).expect("positive")),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
busy_accounts: Default::default(),
preferred_account: None,
preferred_account_overrides_weight: false,
round_robin_cursor: 0,
@@ -17,9 +17,9 @@ use super::{candidate_with_concurrency, context, weighted_candidate};
fn reserved_pool_scoring_and_trace_use_its_capacity_instead_of_the_normal_limit() {
let mut context = context(RotationStrategy::Smart);
context.reserved_concurrency = 5;
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], false).unwrap(),
);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]).unwrap());
let candidates = [
candidate_with_concurrency("acct_large_normal", 2, 100),
candidate_with_concurrency("acct_small_normal", 1, 1),
@@ -44,14 +44,14 @@ fn smart_config_rejects_invalid_weights_and_incomplete_wire_values() {
for dimension in 0..6 {
let mut weights = [1.0; 6];
weights[dimension] = value;
assert!(SmartSchedulingConfig::new(weights, false).is_err());
assert!(SmartSchedulingConfig::new(weights).is_err());
}
}
assert!(SmartSchedulingConfig::new([0.0; 6], true).is_err());
assert!(SmartSchedulingConfig::new([0.0; 6]).is_err());
for value in [json!({}), json!(null), json!({"loadWeight": 1.0})] {
assert!(serde_json::from_value::<SmartSchedulingConfig>(value).is_err());
}
let config = SmartSchedulingConfig::new([0.0, 0.1, 9.9, 10.0, 1.2, 2.3], true).unwrap();
let config = SmartSchedulingConfig::new([0.0, 0.1, 9.9, 10.0, 1.2, 2.3]).unwrap();
let encoded = serde_json::to_value(config).unwrap();
assert_eq!(
serde_json::from_value::<SmartSchedulingConfig>(encoded).unwrap(),
@@ -62,9 +62,9 @@ fn smart_config_rejects_invalid_weights_and_incomplete_wire_values() {
#[test]
fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0], false).unwrap(),
);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0]).unwrap());
let mut candidates = [
candidate_with_concurrency("acct_unknown", 0, 2),
candidate_with_concurrency("acct_past", 0, 2),
@@ -137,7 +137,7 @@ fn queue_weight_balances_live_pressure_and_health_only_when_enabled_for_smart()
candidates[1].signals.failure_rate_basis_points = Some(6000);
let mut context = context(strategy);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 1.0, 0.0, 0.0, queue_weight], false).unwrap(),
SmartSchedulingConfig::new([0.0, 0.0, 1.0, 0.0, 0.0, queue_weight]).unwrap(),
);
let queue = ConcurrencyWaitQueue::default();
let head = queue
@@ -197,7 +197,7 @@ fn queue_weight_does_not_widen_immediate_selection_tolerance() {
context.round_robin_cursor = 1;
for queue_weight in [0.0, 10.0] {
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 1.0, 0.0, 0.0, 0.0, queue_weight], false).unwrap(),
SmartSchedulingConfig::new([0.0, 1.0, 0.0, 0.0, 0.0, queue_weight]).unwrap(),
);
assert_eq!(
AccountSelector
@@ -210,9 +210,9 @@ fn queue_weight_does_not_widen_immediate_selection_tolerance() {
"acct_a"
);
}
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 0.0, 1.0], false).unwrap(),
);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).unwrap());
// 仅开启排队系数时,立即可用的同权重账号按既有游标轮换
assert_eq!(
AccountSelector
@@ -244,7 +244,7 @@ fn each_smart_weight_controls_its_signal_without_relaxing_capacity() {
let mut context = context(RotationStrategy::Smart);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new(weights, false).unwrap());
.with_smart_scheduling(SmartSchedulingConfig::new(weights).unwrap());
let selected = AccountSelector.select(&candidates, &context).unwrap();
assert_eq!(
selected.candidate().account.id().as_str(),
@@ -283,8 +283,7 @@ fn proportional_coefficients_preserve_near_best_rotation_and_diagnostics() {
for scale in [1.0, 2.0, 5.0] {
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([scale, 0.8 * scale, scale, 0.5 * scale, 0.0, 0.0], false)
.unwrap(),
SmartSchedulingConfig::new([scale, 0.8 * scale, scale, 0.5 * scale, 0.0, 0.0]).unwrap(),
);
let mut ids = Vec::new();
for cursor in 0..4 {
@@ -313,64 +312,40 @@ fn proportional_coefficients_preserve_near_best_rotation_and_diagnostics() {
}
#[test]
fn switchback_is_opt_in_and_only_affects_smart_soft_affinity() {
let mut candidates = [
weighted_candidate("acct_affinity", 10, 0),
weighted_candidate("acct_primary", 100, 0),
];
fn higher_weight_recovery_preserves_usable_soft_affinity() {
for strategy in [
RotationStrategy::Smart,
RotationStrategy::Sticky,
RotationStrategy::RoundRobin,
RotationStrategy::QuotaResetPriority,
] {
let mut candidates = [
weighted_candidate("acct_affinity", 10, 0),
weighted_candidate("acct_primary", 100, 3),
];
let mut context = context(strategy);
context.preferred_account = Some(candidates[0].account.id().clone());
context.preferred_account_overrides_weight = true;
for enabled in [false, true] {
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], enabled).unwrap(),
);
for primary_in_flight in [3, 0] {
candidates[1].signals.in_flight = primary_in_flight;
let selected = AccountSelector.select(&candidates, &context).unwrap();
let should_switch = enabled && strategy == RotationStrategy::Smart;
assert_eq!(
selected.candidate().account.id(),
candidates[usize::from(should_switch)].account.id()
candidates[0].account.id()
);
assert_eq!(selected.preferred(), PreferredAccountSelection::Hit);
}
context
.excluded_accounts
.insert(candidates[0].account.id().clone());
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id(),
candidates[1].account.id()
);
}
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true).unwrap(),
);
context.preferred_account_overrides_weight = true;
context.preferred_account = Some(candidates[0].account.id().clone());
candidates[1].signals.in_flight = 3;
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
candidates[1].signals.in_flight = 0;
context
.excluded_accounts
.insert(candidates[1].account.id().clone());
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
context.excluded_accounts.clear();
candidates[0] = weighted_candidate("acct_affinity", 100, 2);
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
}
@@ -66,6 +66,8 @@ fn provider_failure_trace_keeps_classification_when_diagnostic_is_truncated() {
assert_eq!(data["diagnostic"]["stage"], "prepare");
assert_eq!(data["diagnostic"]["code"], "account_proxy_invalid");
assert_eq!(data["diagnostic"]["truncated"], true);
assert_eq!(data["recovery"]["replaySafe"], false);
assert_eq!(data["recovery"]["retryProhibited"], false);
assert!(!snapshot.to_string().contains("PRIVATE_RAW"));
}
@@ -215,6 +215,19 @@ fn client_retry_should_skip_repeated_budget_and_clear_before_completed_delivery(
let mut session = retry_session(&coordinator, id, operation.clone(), "key_client_1");
assert!(block_on(session.collect_uncommitted()).is_err());
assert_eq!(provider.contexts.lock().unwrap().len(), expected);
let trace = session.trace().snapshot().unwrap();
let decision = trace["events"]
.as_array()
.unwrap()
.iter()
.rev()
.find(|event| event["stage"] == "retry.decided")
.unwrap();
assert_eq!(decision["data"]["clientRetrySuppressed"], id == "req_retry");
assert_eq!(decision["data"]["providerProvedReplaySafe"], true);
assert_eq!(decision["data"]["retryProhibited"], false);
assert_eq!(decision["data"]["replayBarrierReached"], false);
assert_eq!(decision["data"]["retryable"], false);
}
let mut recovered = retry_session(
&coordinator,
@@ -4197,7 +4210,7 @@ fn terminal_rejection_should_prevent_not_sent_rotation_and_explicit_transport_re
error = error.with_pre_delivery_transport_fallback();
}
assert!(error.stable_snapshot().retry_is_prohibited());
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
let (coordinator, store, provider) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(error)],
}]);
@@ -4215,6 +4228,27 @@ fn terminal_rejection_should_prevent_not_sent_rotation_and_explicit_transport_re
Err(EngineError::Provider(_))
));
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
let state = store.state.lock().unwrap();
let trace: Value = serde_json::from_str(
state.finalizations[0]
.diagnostic_trace_json
.as_deref()
.unwrap(),
)
.unwrap();
let events = trace["events"].as_array().unwrap();
let failed = events
.iter()
.find(|event| event["stage"] == "attempt.failed")
.unwrap();
assert_eq!(failed["data"]["recovery"]["retryProhibited"], true);
let decision = events
.iter()
.find(|event| event["stage"] == "retry.decided")
.unwrap();
assert_eq!(decision["data"]["phase"], "execute");
assert_eq!(decision["data"]["retryable"], false);
assert_eq!(decision["data"]["retryProhibited"], true);
}
}
@@ -5021,6 +5055,13 @@ fn no_eligible_account_before_stream_records_failure_without_fabricating_an_atte
.iter()
.any(|event| event["stage"] == "attempt.failed")
);
let decision = events
.iter()
.find(|event| event["stage"] == "retry.decided")
.unwrap();
assert_eq!(decision["data"]["phase"], "prepare");
assert_eq!(decision["data"]["retryable"], false);
assert_eq!(decision["data"]["requestSendState"], "not_sent");
assert_eq!(events.last().unwrap()["stage"], "request.finished");
}
}
@@ -6404,6 +6445,29 @@ fn configured_rotation_budget_limits_actual_account_changes() {
.openai_session_affinity_ttl()
== Duration::from_secs(168 * 3600)
}));
let state = store.state.lock().unwrap();
let trace: Value = serde_json::from_str(
state.finalizations[0]
.diagnostic_trace_json
.as_deref()
.unwrap(),
)
.unwrap();
let routed = trace["events"]
.as_array()
.unwrap()
.iter()
.find(|event| event["stage"] == "request.routed")
.unwrap();
assert_eq!(
routed["data"]["accountPolicy"]["maxAccountRotations"],
budget
);
assert_eq!(
routed["data"]["accountPolicy"]["sessionAffinityTtlMs"],
168 * 3600 * 1000_u64
);
drop(state);
assert_eq!(
store.state.lock().unwrap().finalizations[0].attempt_count,
budget + 1
@@ -1966,10 +1966,11 @@ impl Provider for PolicySelectingProvider {
first_output_latency_ms: None,
},
}];
let selection = AccountSelectionContext {
let mut selection = AccountSelectionContext {
policy: attempt.account_selection_policy(),
now: SystemTime::now(),
excluded_accounts: attempt.excluded_accounts().clone(),
busy_accounts: Default::default(),
preferred_account: attempt.required_account().cloned(),
preferred_account_overrides_weight: true,
round_robin_cursor: 0,
@@ -1985,14 +1986,16 @@ impl Provider for PolicySelectingProvider {
.upstream_model()
.map(UpstreamModelId::as_str),
&candidates,
&selection,
&mut selection,
)
.await
{
Err(AccountPolicyError::Rejected) => Err(ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::NotSent,
)),
Err(AccountPolicyError::Rejected | AccountPolicyError::Filtered) => {
Err(ProviderError::new(
ProviderErrorKind::RequestPolicyDenied,
UpstreamSendState::NotSent,
))
}
Err(AccountPolicyError::Fault | AccountPolicyError::StaleCandidate) => Err(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent),
),
@@ -3951,6 +3954,16 @@ fn start_snapshot_with_policy(
enabled: bool,
limits: RateLimits,
fast_mode: FastMode,
) -> RuntimeSnapshot {
start_snapshot_with_image_policy(revision, enabled, limits, fast_mode, true)
}
fn start_snapshot_with_image_policy(
revision: u64,
enabled: bool,
limits: RateLimits,
fast_mode: FastMode,
allow_image_generation: bool,
) -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
@@ -3971,7 +3984,8 @@ fn start_snapshot_with_policy(
account_scope(&provider, "acct_start")
.as_ref()
.clone()
.with_fast_mode(fast_mode),
.with_fast_mode(fast_mode)
.with_image_generation(allow_image_generation),
),
enabled,
limits,
@@ -4895,3 +4909,76 @@ fn child_settings_recompute_key_scope_without_inheriting_parent_defaults() {
);
});
}
#[test]
fn image_policy_stays_frozen_for_prepared_requests_and_refreshes_for_the_next_turn() {
block_on(async {
let snapshot = |revision, allowed| {
start_snapshot_with_image_policy(
revision,
true,
RateLimits::unlimited(),
FastMode::Default,
allowed,
)
};
let snapshots = RuntimeSnapshotHandle::new(snapshot(1, false));
let admissions = Arc::new(Admissions::default());
let service = DefaultExecutionService::new(
snapshots.clone(),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
admissions.clone(),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
Arc::new(crate::RecordingDiagnostics::default()),
);
let client = service.authenticate("sk_start_test").unwrap();
let prepared = service.prepare_execution(client.clone()).await.unwrap();
// 关闭生图不影响文本请求的正常准入
let text = service
.start(request(&service, ClientTransport::HttpJson))
.await
.unwrap();
text.session.detach_finalize().await;
assert_eq!(admissions.limits.lock().unwrap().len(), 1);
snapshots.publish(snapshot(2, true));
assert!(
!prepared
.client()
.policy()
.account_scope()
.allows_image_generation()
);
let next = service.prepare_execution(client).await.unwrap();
assert!(
next.client()
.policy()
.account_scope()
.allows_image_generation()
);
let input = request(&service, ClientTransport::WebSocket);
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
json!({"model":"gpt-start","tools":[{"type":"image_generation"}]})
.as_object()
.unwrap()
.clone(),
)
.unwrap(),
));
let result = service
.start_prepared(
prepared,
gateway_core::engine::execution::PreparedExecutionRequest {
public_model: input.public_model,
operation,
metadata: input.metadata,
},
)
.await;
assert!(matches!(result, Err(error) if error.kind() == GatewayErrorKind::PolicyDenied));
assert_eq!(admissions.limits.lock().unwrap().len(), 1);
});
}
@@ -661,3 +661,97 @@ fn attempt_middleware_preserves_multi_value_and_authentication_headers() {
assert_eq!(calls.load(Ordering::SeqCst), 1);
});
}
#[derive(Debug)]
struct AddingImageTool;
impl MiddlewarePlan for AddingImageTool {
fn handle(
&self,
_: MiddlewareContext,
request: MiddlewareRequest,
next: MiddlewareNext,
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
Box::pin(async move {
let (protocol, headers, body) = request.into_parts();
let mut body: Value = serde_json::from_slice(&body).unwrap();
body["tools"] = json!([{"type":"image_generation"}]);
next.run(MiddlewareRequest::new(
protocol,
headers,
serde_json::to_vec(&body).unwrap().into(),
))
.await
})
}
}
#[test]
fn attempt_middleware_cannot_add_image_generation_to_a_denied_scope() {
use gateway_core::account::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::{AccountAttemptContext, AttemptContext, RequestAttemptContext};
use gateway_core::error::ProviderErrorKind;
use gateway_core::routing::{ClientRoutingScope, FrozenAccountScope, RuntimeAccountDirectory};
use gateway_core::upstream::UpstreamSendState;
for allowed in [false, true] {
let scope = FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(Default::default())),
ClientRoutingScope::all_accounts(),
)
.with_image_generation(allowed);
let request = RequestAttemptContext::new(
ModelRequestId::new("req_image_policy").unwrap(),
ClientApiKeyId::new("key_image_policy").unwrap(),
)
.with_middleware(
Some(FrozenMiddlewarePlan::new(
Arc::new(AddingImageTool),
ExtensionSetReference::new(
ExtensionSetId::new("image-policy-test".to_owned()).unwrap(),
Arc::new(ReadyGeneration),
),
)),
Arc::from([]),
"/v1/responses".into(),
ClientTransport::HttpJson,
);
let attempt = AttemptContext::new(
request,
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(1).unwrap()),
AccountAttemptContext::new(Default::default(), None, None)
.with_account_scope(Arc::new(scope)),
None,
CancellationToken::new(),
);
let calls = Arc::new(AtomicUsize::new(0));
let count = calls.clone();
let result = futures::executor::block_on(attempt.execute_middleware(
body_operation(&json!({"model":"gpt-test","input":"hello"})),
ProviderKind::new("openai").unwrap(),
Some("gpt-test".into()),
ProviderAccountId::new("acct_middleware").unwrap(),
Box::new(move |operation, _| {
Box::pin(async move {
assert!(operation.image_generation_requested());
count.fetch_add(1, Ordering::SeqCst);
Ok(empty_provider_stream())
})
}),
));
assert_eq!(calls.load(Ordering::SeqCst), usize::from(allowed));
if allowed {
assert!(result.is_ok());
} else {
let error = match result {
Err(error) => error,
Ok(_) => panic!("image tool must be denied"),
};
assert_eq!(error.kind(), ProviderErrorKind::RequestPolicyDenied);
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
assert!(error.retry_is_prohibited());
}
}
}
@@ -23,7 +23,7 @@ use gateway_core::{
};
#[derive(Debug)]
struct FixedScheduler(Option<&'static str>);
struct FixedScheduler(AccountScheduleDecision);
impl RequestPolicyPlan for FixedScheduler {
fn route_model(
@@ -37,9 +37,7 @@ impl RequestPolicyPlan for FixedScheduler {
&self,
_: AccountScheduleInput,
) -> BoxFuture<'static, Result<AccountScheduleDecision, RequestPolicyFault>> {
let decision = self.0.map_or(AccountScheduleDecision::Delegate, |id| {
AccountScheduleDecision::Pick(account_id(id))
});
let decision = self.0.clone();
Box::pin(async move { Ok(decision) })
}
}
@@ -87,8 +85,14 @@ fn candidate(id: &str) -> AccountCandidate {
}
fn policy(choice: Option<&'static str>) -> RequestPolicyContext {
policy_decision(choice.map_or(AccountScheduleDecision::Delegate, |id| {
AccountScheduleDecision::Pick(account_id(id))
}))
}
fn policy_decision(decision: AccountScheduleDecision) -> RequestPolicyContext {
RequestPolicyContext::new(
Arc::new(FixedScheduler(choice)),
Arc::new(FixedScheduler(decision)),
ExtensionSetReference::new(
ExtensionSetId::new("scheduler-test".into()).unwrap(),
Arc::new(Lease),
@@ -104,6 +108,7 @@ fn context() -> AccountSelectionContext {
policy: AccountSelectionPolicy::new(RotationStrategy::Sticky, NonZeroU32::new(2).unwrap()),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
busy_accounts: Default::default(),
preferred_account: Some(account_id("acct_a")),
preferred_account_overrides_weight: true,
round_robin_cursor: 0,
@@ -123,7 +128,7 @@ fn scheduler_can_override_soft_affinity_and_reports_the_override() {
&ProviderKind::new("openai").unwrap(),
Some("demo-model"),
&candidates,
&context(),
&mut context(),
)
.await
.unwrap()
@@ -148,7 +153,7 @@ fn delegated_scheduler_keeps_builtin_affinity() {
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&context(),
&mut context(),
)
.await
.unwrap()
@@ -164,7 +169,7 @@ fn delegated_scheduler_keeps_builtin_affinity() {
}
#[test]
fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry() {
fn explicit_plugin_choices_override_native_soft_affinity() {
block_on(async {
let mut candidates = [candidate("acct_a"), candidate("acct_b")];
candidates[0].account = candidates[0]
@@ -175,10 +180,9 @@ fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry
selection_context.policy =
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[1.0, 0.8, 1.0, 0.5, 0.0, 0.0],
true,
)
gateway_core::account::SmartSchedulingConfig::new([
1.0, 0.8, 1.0, 0.5, 0.0, 0.0,
])
.unwrap(),
);
for chosen in [Some("acct_a"), Some("acct_b"), None] {
@@ -188,21 +192,19 @@ fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&selection_context,
&mut selection_context,
)
.await
.unwrap()
.unwrap();
assert_eq!(
selected.candidate().account.id().as_str(),
chosen.unwrap_or("acct_b")
chosen.unwrap_or("acct_a")
);
let expected = match chosen {
Some("acct_a") => PreferredAccountSelection::Hit,
Some(_) => PreferredAccountSelection::OverriddenByPolicy,
None => PreferredAccountSelection::Blocked(
gateway_core::account::AccountSchedulingBlocker::LowerWeight,
),
None => PreferredAccountSelection::Hit,
};
assert_eq!(selected.preferred(), expected);
}
@@ -221,7 +223,181 @@ fn scheduler_cannot_select_outside_the_hard_bound_candidate_pool() {
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&context,
&mut context,
)
.await;
assert!(matches!(result, Err(AccountPolicyError::Fault)));
});
}
#[test]
fn filtered_delegation_preserves_native_weights_and_affinity_on_remaining_accounts() {
block_on(async {
let mut candidates = [
candidate("acct_a"),
candidate("acct_b"),
candidate("acct_c"),
];
candidates[1].account = candidates[1]
.account
.clone()
.with_scheduling(None, AccountWeight::new(10).unwrap());
let policy = policy_decision(AccountScheduleDecision::DelegateFiltered(BTreeSet::from([
account_id("acct_a"),
])));
for strategy in [RotationStrategy::Smart, RotationStrategy::RoundRobin] {
let mut context = context();
context.policy = AccountSelectionPolicy::new(strategy, NonZeroU32::new(2).unwrap());
let selected = policy
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context,
)
.await
.unwrap()
.unwrap();
assert_eq!(selected.candidate().account.id(), &account_id("acct_c"));
assert!(!selected.is_policy_choice());
assert_eq!(
gateway_core::account::AccountSelector.wait_candidates(&candidates, &context),
[account_id("acct_c"), account_id("acct_b")]
);
assert_eq!(
gateway_core::account::AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id(),
&account_id("acct_c")
);
}
});
}
#[test]
fn filtered_delegation_covers_busy_candidates_and_does_not_persist_between_selections() {
block_on(async {
let mut candidates = [candidate("acct_a"), candidate("acct_b")];
candidates[0].signals.in_flight = 2;
let mut selection = context();
selection.busy_accounts.insert(account_id("acct_b"));
let filtered =
policy_decision(AccountScheduleDecision::DelegateFiltered(BTreeSet::from([
account_id("acct_a"),
])));
assert!(
filtered
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut selection
)
.await
.unwrap()
.is_none()
);
assert_eq!(
gateway_core::account::AccountSelector.wait_candidates(&candidates, &selection),
[account_id("acct_b")]
);
candidates[0].signals.in_flight = 0;
let recovered = policy(None)
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context(),
)
.await
.unwrap()
.unwrap();
assert_eq!(recovered.candidate().account.id(), &account_id("acct_a"));
});
}
#[test]
fn filtering_all_candidates_is_terminal_even_when_the_only_candidate_is_busy_or_pinned() {
block_on(async {
for busy in [false, true] {
let candidates = [candidate("acct_a"), candidate("acct_b")];
let mut context = context();
context.excluded_accounts.insert(account_id("acct_b"));
if busy {
context.busy_accounts.insert(account_id("acct_a"));
}
let result =
policy_decision(AccountScheduleDecision::DelegateFiltered(BTreeSet::from([
account_id("acct_a"),
])))
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context,
)
.await;
assert!(matches!(result, Err(AccountPolicyError::Filtered)));
}
});
}
#[test]
fn filtering_ids_outside_the_host_candidate_pool_is_a_policy_fault() {
block_on(async {
let candidates = [candidate("acct_a"), candidate("acct_b")];
for id in ["acct_b", "acct_unknown"] {
let mut context = context();
context.excluded_accounts.insert(account_id("acct_b"));
let result =
policy_decision(AccountScheduleDecision::DelegateFiltered(BTreeSet::from([
account_id(id),
])))
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context,
)
.await;
assert!(matches!(result, Err(AccountPolicyError::Fault)));
}
});
}
#[test]
fn empty_filter_preserves_native_affinity_and_busy_candidates_cannot_be_picked() {
block_on(async {
let candidates = [candidate("acct_a"), candidate("acct_b")];
let selected = policy_decision(AccountScheduleDecision::DelegateFiltered(BTreeSet::new()))
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context(),
)
.await
.unwrap()
.unwrap();
assert_eq!(selected.preferred(), PreferredAccountSelection::Hit);
let mut context = context();
context.busy_accounts.insert(account_id("acct_a"));
let result = policy(Some("acct_a"))
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&mut context,
)
.await;
assert!(matches!(result, Err(AccountPolicyError::Fault)));
@@ -156,6 +156,55 @@ fn prompt_cache_and_image_generation_should_be_read_from_raw_body() {
assert!(!format!("{request:?}").contains("private-cache-route"));
}
#[test]
fn image_generation_detection_distinguishes_tools_from_image_inputs_and_passive_namespaces() {
for body in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
json!({"tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]}}),
json!({"input":[{"type":"additional_tools","tools":[{"type":"image_generation"}]}]}),
] {
assert!(
generate(body.clone()).image_generation_requested(),
"{body}"
);
}
for body in [
json!({"input":"Please generate an image"}),
json!({"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}),
json!({"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}),
json!({"tools":[{"type":"function","name":"image_generation"}]}),
] {
assert!(
!generate(body.clone()).image_generation_requested(),
"{body}"
);
}
let private = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"private",
json!({"tools":[{"type":"image_generation"}]})
.as_object()
.unwrap()
.clone(),
)
.unwrap(),
);
assert!(!private.image_generation_requested());
}
#[test]
fn image_generation_fact_survives_cross_protocol_translation() {
let operation = Operation::Generate(generate(json!({"tools":[{"type":"image_generation"}]})));
let translated = operation
.replace_middleware_wire(
"openai_chat_completions",
Bytes::from_static(br#"{"model":"gpt-test","messages":[]}"#),
)
.unwrap();
assert!(translated.image_generation_requested());
}
#[test]
fn protocol_payload_context_should_be_opaque_and_separate_from_wire_body() {
let secret = "connection-only-value";
@@ -443,7 +443,7 @@ fn compiled_plans_keep_their_scheduling_config_after_a_new_snapshot_is_built() {
use gateway_core::account::SmartSchedulingConfig;
let configs = [
SmartSchedulingConfig::default(),
SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3], true).unwrap(),
SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3]).unwrap(),
];
let mut plans = Vec::new();
for (index, config) in configs.into_iter().enumerate() {
@@ -1305,6 +1305,79 @@ fn fast_mode_merges_only_bound_groups_without_changing_account_scope() {
}
}
#[test]
fn image_generation_merges_bound_groups_with_deny_precedence() {
use gateway_core::account::ProviderAccountId;
use gateway_core::routing::AccountGroupId;
for (allowed, other_allowed) in [(true, true), (true, false), (false, true), (false, false)] {
for group_enabled in [false, true] {
for bound in [false, true] {
let group_id = AccountGroupId::new("grp_00000000000000000000000000000001").unwrap();
let open_group_id =
AccountGroupId::new("grp_00000000000000000000000000000002").unwrap();
let account_id = ProviderAccountId::new("acct_fast_policy").unwrap();
let facts = SnapshotFacts::new(
revision(1),
revision(1),
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_fast_policy").unwrap(),
PlaintextClientApiKey::new("sk_fast_policy").unwrap(),
if bound {
vec![group_id.clone(), open_group_id.clone()]
} else {
Vec::new()
},
RateLimits::unlimited(),
)],
vec![
SnapshotAccountGroupFacts::new(
group_id.clone(),
"Restricted".to_owned(),
group_enabled,
)
.with_image_generation(allowed),
SnapshotAccountGroupFacts::new(
open_group_id.clone(),
"Open".to_owned(),
true,
)
.with_image_generation(other_allowed),
],
vec![SnapshotProviderAccountFacts::new(
account_id.clone(),
"alpha",
)],
vec![
SnapshotAccountGroupMemberFacts::new(group_id, account_id.clone()),
SnapshotAccountGroupMemberFacts::new(open_group_id, account_id.clone()),
],
);
let snapshot = block_on(
RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts))),
Arc::new(TestCatalog::Unavailable),
)
.compile(),
)
.unwrap();
let scope = snapshot
.client_policies()
.next()
.unwrap()
.account_scope()
.clone();
assert!(scope.allows(&account_id));
assert_eq!(
scope.allows_image_generation(),
!bound || (allowed && other_allowed),
"allowed={allowed}, other={other_allowed}, enabled={group_enabled}, bound={bound}"
);
}
}
}
}
#[test]
fn key_profiles_replace_whole_global_choice_and_previous_snapshot_stays_frozen() {
use gateway_core::account::OpaqueProviderData;
@@ -4,7 +4,7 @@
"protocol_versions": [2],
"capabilities": [
{ "capability": "frontend_authentication", "versions": [1] },
{ "capability": "scheduler", "versions": [1] },
{ "capability": "scheduler", "versions": [2] },
{ "capability": "model_router", "versions": [1] },
{ "capability": "model_catalog", "versions": [1] },
{ "capability": "retry_policy", "versions": [1] },
@@ -196,6 +196,7 @@ async fn schedule_with_entry(
.iter()
.map(|candidate| WireAccountScheduleCandidate {
account_id: candidate.account_id().as_str().to_owned(),
available: candidate.available(),
weight: candidate.weight(),
in_flight: candidate.in_flight(),
maximum_concurrency: candidate.maximum_concurrency(),
@@ -245,11 +246,27 @@ async fn schedule_with_entry(
match decision {
WireAccountScheduleDecision::Delegate => Ok(AccountScheduleDecision::Delegate),
WireAccountScheduleDecision::Reject => Ok(AccountScheduleDecision::Reject),
WireAccountScheduleDecision::DelegateFiltered {
exclude_account_ids,
} => {
let candidate_ids = candidates
.iter()
.map(|candidate| candidate.account_id())
.collect::<std::collections::HashSet<_>>();
let excluded = exclude_account_ids
.into_iter()
.map(|id| {
let id = ProviderAccountId::new(id).map_err(|_| ())?;
candidate_ids.contains(&id).then_some(id).ok_or(())
})
.collect::<Result<_, _>>()?;
Ok(AccountScheduleDecision::DelegateFiltered(excluded))
}
WireAccountScheduleDecision::Pick { account_id } => {
let account_id = ProviderAccountId::new(account_id).map_err(|_| ())?;
if !candidates
.iter()
.any(|candidate| candidate.account_id() == &account_id)
.any(|candidate| candidate.account_id() == &account_id && candidate.available())
{
return Err(());
}
@@ -81,6 +81,7 @@ impl PluginCallbackPorts {
pub(crate) struct PluginCallbacks {
services: Arc<services::ServicePorts>,
service_actor: gateway_admin::model::usage_data::UsageActor,
accounts: Arc<accounts::PluginAccounts>,
resources: Arc<resources::PluginResources>,
data: Arc<data::PluginData>,
@@ -108,6 +109,11 @@ impl PluginCallbacks {
) -> Result<Self, AdminError> {
Ok(Self {
services: ports.services,
service_actor: gateway_admin::model::usage_data::UsageActor::Plugin {
publisher: manifest.publisher.clone(),
name: manifest.name.clone(),
instance_id: instance.id.clone(),
},
resources: Arc::new(resources::PluginResources::new(instance, ports.resources)),
data: Arc::new(data::PluginData::new(
ports.accounts.clone(),
@@ -386,6 +392,7 @@ impl CallbackHandler for PluginCallbacks {
let models = self.models.clone();
let affinity = self.affinity.clone();
let services = self.services.clone();
let service_actor = self.service_actor.clone();
Box::pin(async move {
let call = call.ok_or_else(denied)?;
let scope = &call.scope;
@@ -432,7 +439,14 @@ impl CallbackHandler for PluginCallbacks {
}
if method == gateway_plugin_sdk::call::services::CALL_METHOD {
return services
.call(&context, scope, call.cancellation.clone(), params, &payload)
.call(
&context,
scope,
call.cancellation.clone(),
service_actor,
params,
&payload,
)
.await;
}
if is_private_state {
@@ -73,6 +73,7 @@ impl PluginResources {
&self.owner,
request.resource_key,
CreateAccountGroup {
allow_image_generation: true,
fast_mode: gateway_core::account::FastMode::Default,
name: request.name,
description: request.description,
@@ -49,6 +49,7 @@ impl ServicePorts {
context: &CallContext,
scope: &CallbackScope,
cancellation: gateway_core::lifecycle::CancellationToken,
actor: gateway_admin::model::usage_data::UsageActor,
params: serde_json::Value,
payload: &[u8],
) -> Result<RpcReply, PluginFault> {
@@ -84,7 +85,11 @@ impl ServicePorts {
gateway_admin::public_service::Plan::Frozen(Some(origin.plan.clone()))
}),
};
let result = gateway_admin::public_service::scope(parent, call).await;
let result = gateway_admin::public_service::scope_actor(
actor,
gateway_admin::public_service::scope(parent, call),
)
.await;
encode(result)
}
}
@@ -20,6 +20,7 @@ use std::{sync::Arc, time::Duration};
fn group(name: &str) -> CreateAccountGroup {
CreateAccountGroup {
allow_image_generation: true,
name: name.into(),
description: None,
fast_mode: FastMode::Default,
@@ -5,6 +5,7 @@ mod facts;
mod http;
mod lifecycle;
mod mounts;
mod scheduling;
mod service;
mod websocket;
@@ -556,6 +557,7 @@ fn selection_context() -> AccountSelectionContext {
),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
busy_accounts: Default::default(),
preferred_account: None,
preferred_account_overrides_weight: true,
round_robin_cursor: 0,
@@ -582,14 +584,14 @@ async fn scheduler_cannot_pick_outside_the_host_projected_priority_tier() {
let generation = prepare(&runtime).await;
let context = policy_context(&runtime, generation.clone(), "req_schedule");
let candidates = [candidate("acct_high", 100), candidate("acct_low", 10)];
let selection = selection_context();
let mut selection = selection_context();
let selected = context
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("example").unwrap(),
Some("upstream-a"),
&candidates,
&selection,
&mut selection,
)
.await
.unwrap()
@@ -0,0 +1,69 @@
//! 验证真实插件进程的候选过滤、繁忙投影与非法决定边界
use super::*;
use gateway_core::engine::policy::AccountPolicyError;
#[tokio::test]
async fn scheduler_filter_contract_covers_busy_candidates_and_rejects_out_of_scope_ids() {
for (decision, expected) in [
(
serde_json::json!({"decision":"delegate_filtered","exclude_account_ids":["acct_high", "acct_high"]}),
"acct_low",
),
(
serde_json::json!({"decision":"delegate_filtered","exclude_account_ids":["acct_high", "acct_low"]}),
"filtered",
),
(
serde_json::json!({"decision":"delegate_filtered","exclude_account_ids":["acct_unknown"]}),
"fault",
),
(
serde_json::json!({"decision":"pick","account_id":"acct_high"}),
"fault",
),
] {
let marker = tempfile::NamedTempFile::new().unwrap();
let (cache, runtime) = setup(vec![InstanceFixture {
id: "scheduler",
configuration: serde_json::json!({"schedule_decision":decision, "schedule_marker": marker.path()}),
bindings: vec![binding(SCHEDULER_CONTRIBUTION, "scheduling", 0, PluginFailurePolicy::Reject)],
}]).await;
let generation = prepare(&runtime).await;
let policy = policy_context(&runtime, generation.clone(), "req_filter");
let candidates = [candidate("acct_high", 100), candidate("acct_low", 10)];
let mut selection = selection_context();
selection
.busy_accounts
.insert(candidates[0].account.id().clone());
selection.preferred_account = Some(candidates[0].account.id().clone());
let result = policy
.select_account(
NonZeroU32::MIN,
&ProviderKind::new("example").unwrap(),
None,
&candidates,
&mut selection,
)
.await;
match expected {
"filtered" => assert!(matches!(result, Err(AccountPolicyError::Filtered))),
"fault" => assert!(matches!(result, Err(AccountPolicyError::Fault))),
id => assert_eq!(
result.unwrap().unwrap().candidate().account.id().as_str(),
id
),
}
let request: serde_json::Value =
serde_json::from_str(std::fs::read_to_string(marker.path()).unwrap().trim()).unwrap();
assert_eq!(
request["request"]["candidates"].as_array().unwrap().len(),
2
);
assert_eq!(request["request"]["candidates"][0]["available"], false);
assert_eq!(request["request"]["candidates"][1]["available"], true);
drop(policy);
drop(generation);
super::super::wait_until_empty(cache.path()).await;
}
}
@@ -459,3 +459,46 @@ async fn service_chain_preserves_onion_order_parent_scope_and_native_results_wit
runtime.shutdown().await;
drop(cache);
}
#[tokio::test]
async fn usage_sdk_calls_preserve_exact_counts_and_write_actor() {
use gateway_admin::model::usage_data::*;
let Some(environment) = Environment::create().await else {
return;
};
let (runtime, core) = environment
.plugin(json!({"service":true,"mode":"usage_data"}))
.await;
let admin = environment.bind_admin_accounts(&runtime, &core).await;
let services = admin.services();
runtime.bind_services(&services.public_services()).unwrap();
assert_eq!(
services
.public_services()
.call("settings.admin_api_key_exists", json!(null))
.await
.unwrap(),
json!(true)
);
let row = services
.usage_data()
.get_record(GetUsageRecord {
id: "sdk-usage".into(),
})
.await
.unwrap();
assert_eq!(
row.data.tokens.total_tokens,
Some(UsageCount(9_007_199_254_740_993))
);
assert!(matches!(row.data.cost, UsageCost::Known { .. }));
let actors = environment.audit_actor_refs("usage.put").await;
assert!(!actors.is_empty());
assert!(
actors
.iter()
.all(|actor| actor.starts_with("plugin:test/example:"))
);
runtime.shutdown().await;
}
@@ -65,7 +65,7 @@ impl Environment {
.await
.unwrap();
sqlx::query(sqlx::AssertSqlSafe(format!(
"insert into {}.client_key_budget_windows
"insert into {}.client_api_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
@@ -127,6 +127,7 @@ impl Environment {
.account_groups()
.create_account_group(
NewAccountGroup {
allow_image_generation: true,
fast_mode: FastMode::Default,
id: id.clone(),
name: format!("fixture {}", id.as_str()),
@@ -720,6 +721,17 @@ impl Environment {
self.admin.close().await;
}
pub async fn audit_actor_refs(&self, action: &str) -> Vec<String> {
sqlx::query_scalar(sqlx::AssertSqlSafe(format!(
"select actor_ref from {}.admin_audit_events where action=$1 order by created_at",
self.schema
)))
.bind(action)
.fetch_all(&self.admin)
.await
.unwrap()
}
pub async fn audit_requests(&self, action: &str) -> Vec<String> {
sqlx::query_scalar(sqlx::AssertSqlSafe(format!(
"select admin_request_id from {}.admin_audit_events where action = $1 order by created_at",
@@ -1,6 +1,7 @@
//! 中间件测试插件子进程入口及异常退出清理辅助
mod http_resources;
mod usage_data;
use gateway_plugin_sdk::{
Capability, ContributionDeclaration, Contributions, Stage,
@@ -179,6 +180,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
));
}
use gateway_plugin_sdk::{call::services::settings, client::ServiceResponse};
if mode == "usage_data" && call.is::<settings::ApiKeyExists>() {
return usage_data::exercise(call).await;
}
if call.is::<settings::ApiKeyExists>() && mode == "recover" {
return ServiceResponse::from_result::<settings::ApiKeyExists>(Ok(true));
}
@@ -0,0 +1,146 @@
//! 真实插件子进程使用独立 SDK 组合明细、聚合和计价基础能力
use gateway_plugin_sdk::{
PluginFault,
call::services::{settings, usage::*},
client::{ServiceCall, ServiceResponse},
};
pub(super) async fn exercise(call: ServiceCall) -> Result<ServiceResponse, PluginFault> {
let call = call.into_typed::<settings::ApiKeyExists>()?;
let data = UsageRecordData {
client_key_ref: "sdk-key".into(),
provider: Some("openai".into()),
account_ref: None,
requested_model: Some("gpt-6-sol".into()),
model: Some("gpt-6-sol".into()),
protocol: "openai".into(),
operation: "responses".into(),
endpoint: "/v1/responses".into(),
client_transport: "http_sse".into(),
service_tier: None,
started_at_ms: 1780272000000,
completed_at_ms: Some(1780272001000),
outcome: UsageOutcome::Succeeded,
tokens: UsageTokens {
input_tokens: Some(UsageCount(100)),
output_tokens: Some(UsageCount(20)),
cached_input_tokens: Some(UsageCount(0)),
cache_write_input_tokens: Some(UsageCount(0)),
image_input_tokens: Some(UsageCount(0)),
image_output_tokens: Some(UsageCount(0)),
total_tokens: Some(UsageCount(9_007_199_254_740_993)),
..Default::default()
},
cost: UsageCost::Unavailable,
};
let write = WriteUsageRecords {
changes: vec![UsageRecordChange::Put {
id: "sdk-usage".into(),
data: Box::new(data.clone()),
}],
};
call.host
.service::<WriteRecords>(write.clone())
.await
.unwrap();
call.host.service::<WriteRecords>(write).await.unwrap();
let row = call
.host
.service::<GetRecord>(GetUsageRecord {
id: "sdk-usage".into(),
})
.await
.unwrap();
assert_eq!(
row.data.tokens.total_tokens,
Some(UsageCount(9_007_199_254_740_993))
);
let filter = UsageRecordFilter {
start_ms: 1780272000000,
end_ms: 1780272002000,
provider: None,
model: None,
account_ref: None,
client_key_ref: None,
};
assert_eq!(
call.host
.service::<ListRecords>(ListUsageRecords {
filter: filter.clone(),
cursor: None,
limit: Some(10)
})
.await
.unwrap()
.items
.len(),
1
);
let quote=call.host.service::<EvaluatePricing>(EvaluateUsagePricing{items:vec![UsagePricingInput{id:"quote".into(),provider:"openai".into(),model:"gpt-6-sol".into(),tokens:data.tokens,service_tier:None,price:UsagePriceSelection::Override{price:Box::new(serde_json::from_value(serde_json::json!({"multiplierBps":10000,"bands":{"standard":{"input":"2","output":"10","cacheRead":"0.5","cacheWrite":"2"}}})).unwrap())}}]}).await.unwrap();
let UsagePricingResult::Estimated { cost, .. } = quote.items[0].clone() else {
panic!("expected quote")
};
call.host
.service::<WriteRecords>(WriteUsageRecords {
changes: vec![UsageRecordChange::Update {
id: "sdk-usage".into(),
tokens: None,
cost: Some(cost),
}],
})
.await
.unwrap();
let identity = UsageAggregateKey {
bucket_start_ms: 1780272000000,
dimensions: UsageDimensions::default(),
key: "sdk-daily".into(),
};
call.host
.service::<WriteAggregates>(WriteUsageAggregates {
changes: vec![UsageAggregateChange::Put {
aggregate: UsageAggregate {
identity: identity.clone(),
counts: std::collections::BTreeMap::from([(
UsageMetric::RequestCount,
UsageDelta(100),
)]),
costs: Default::default(),
},
}],
})
.await
.unwrap();
assert_eq!(
call.host
.service::<ListAggregates>(ListUsageAggregates {
filter,
cursor: None,
limit: None
})
.await
.unwrap()
.items
.len(),
1
);
call.host
.service::<WriteAggregates>(WriteUsageAggregates {
changes: vec![UsageAggregateChange::Delete { identity }],
})
.await
.unwrap();
call.host
.service::<WriteRecords>(WriteUsageRecords {
changes: vec![
UsageRecordChange::Put {
id: "sdk-delete".into(),
data: Box::new(row.data),
},
UsageRecordChange::Delete {
id: "sdk-delete".into(),
},
],
})
.await
.unwrap();
ServiceResponse::from_result::<settings::ApiKeyExists>(Ok(true))
}
@@ -59,7 +59,7 @@ pub fn contribution_for_id(
id: format!("{plugin_id}.{local_id}"),
version: match capability {
Capability::Middleware => 4,
Capability::UpstreamAdapter => 2,
Capability::Scheduler | Capability::UpstreamAdapter => 2,
_ => 1,
},
stages,
@@ -188,6 +188,22 @@ impl NativeAdmin {
#[async_trait]
impl ProviderAdmin for NativeAdmin {
fn evaluate_usage_pricing(
&self,
_: &gateway_admin::ports::provider::ProviderUsagePricingInput,
_: Option<&gateway_core::metering::ModelPriceOverride>,
) -> Result<Option<gateway_core::metering::CalculatedCostBreakdown>, ProviderAdminError> {
use gateway_core::metering::*;
let zero = Money::new(Decimal::ZERO, CurrencyCode::new("USD").unwrap());
let total = Money::new("1.25".parse().unwrap(), CurrencyCode::new("USD").unwrap());
Ok(Some(CalculatedCostBreakdown::new(
CalculatedCostAmounts::new(total, zero, zero, zero, total, total),
CalculatedCostRates::new(total, zero, zero, zero),
None,
100,
)))
}
fn provider_kind(&self) -> &ProviderKind {
&self.0
}
+4 -2
View File
@@ -9,7 +9,7 @@ SDK 处于实验阶段;能力是否可用取决于宿主支持、清单声明
| SDK 包 | `0.1.0` | Rust 开发依赖,独立于网关版本 |
| `manifestVersion` | `2` | 安装清单格式 |
| `package.protocolVersion` | `2` | 宿主与插件的进程通信 |
| `contributes.*.version` | 默认 `1`,`middleware` 为 `4`,`upstream_adapter` 为 `2` | 单项能力合同,见[扩展项声明](docs/manifest.md#扩展项简写) |
| `contributes.*.version` | 默认 `1`,`middleware` 为 `4`,`scheduler` / `upstream_adapter` 为 `2` | 单项能力合同,见[扩展项声明](docs/manifest.md#扩展项简写) |
开发顺序:编写[清单](docs/manifest.md) → 实现[能力处理器](docs/capabilities.md) → 用 [CLI](../../../apps/plugin-cli/README.md)
打包 → [安装、配置与启用](../../../../docs/plugins.md)。完整可运行示例位于独立仓库
@@ -25,7 +25,7 @@ SDK 处于实验阶段;能力是否可用取决于宿主支持、清单声明
| `call::frontend_authentication` | 数据面认证信封、认证器标识与外部 principal 结果 |
| `call::model` | 宿主模型调用的 canonical 事实、原生 wire 与有界事件编解码 |
| `call::middleware` | HTTP、WebSocket 消息、公开服务、模型请求及 attempt 的洋葱调用、单次 next 与惰性正文合同 |
| `call::services` | 公开服务的类型化操作、输入输出及错误合同 |
| `call::services` | 公开服务的类型化操作、输入输出及错误合同;`usage` 提供[用量维护与计价](docs/usage-data.md) |
| `call::upstream_adapter` | 内置 OpenAI / xAI 的[受管上游适配器](docs/upstream-adapters.md)注册、执行、出站与续接合同 |
| `call::policy` | 模型路由、账号调度与重试决策合同 |
| `call::catalog` | 固定 Provider 的模型别名注册合同 |
@@ -113,6 +113,8 @@ CPR_UPDATE_SERVICE_CONTRACT=1 RUST_MIN_STACK=16777216 cargo +1.97.0 test --manif
cargo +1.97.0 fmt --all --manifest-path backend/Cargo.toml
```
用量服务合同由 Admin 与 SDK 分别持有,`sdk_wire_types_match_host_contracts` 检查数值、明细、聚合、查询和计价类型一致性;修改合同需同步两侧类型与 [使用说明](docs/usage-data.md)
SDK 自身的验证命令:
```bash
@@ -403,6 +403,8 @@ restore_and_validate_response(response).await
仅实现中间件时也可使用轻量的 `MiddlewarePlugin`;需要与管理或其他方法组合时使用
`PluginBuilder`,两者复用相同的类型化调用与结果合同
用量明细、聚合维护与费用试算使用 `call::services::usage`,输入输出和批量语义见 [用量维护与计价](usage-data.md)
### 扩展的组合边界
请求处理沿 `request middleware → 原生选路与选号 → attempt middleware → 原生上游或 upstream_adapter` 推进,
@@ -467,6 +469,24 @@ middleware 的 `next` 最多调用一次,重试与额外模型执行通过独
### 路由、调度与观察
`policy.route_model` 和 `policy.schedule_account` 只提出决定,宿主随后复核 Key 规则、模型能力、账号资格和租约。
`scheduler` 必须声明 `version: 2`。`candidates` 已按 Key 范围、模型权限、账号状态和原生续写绑定收窄,
同时包含暂时繁忙但可以等待的账号;`available` 表示本次可以立即尝试租约,不保证租约一定成功
- `Pick { account_id }` 显式选号,只能指定 `available: true` 的候选
- `Delegate` 将完整候选交给宿主现有调度
- `DelegateFiltered { exclude_account_ids }` 排除列出的候选,再委托内置权重、Smart 评分和亲和选择;仅接受本次候选 ID,重复 ID 合并,空列表等同于 `Delegate`
- `Reject` 明确拒绝请求
过滤结果约束本次评估后的亲和重选、容量判断和等待队列。全部候选被排除,或严格跟随的会话账号被排除时,
请求以 `account_policy_filtered` 结束,不回退到被排除账号或伪装成额度耗尽。租约争用、排队唤醒和后续请求重新评估策略,
过滤结果不写入账号状态。插件决定自身阈值、恢复条件和配置含义;宿主只检查合同、候选边界并执行决定。
插件设置使用自己的 `configurationSchema` 或管理页面,不向宿主智能调度弹窗注入字段。
非法输出按绑定的既有故障策略处理,不自动修正插件业务规则
```json
{"decision":"delegate_filtered","exclude_account_ids":["acct_example"]}
```
`observer` v1 通过 `methods::OBSERVE` 接收 `call::observation::Event`,固定绑定 `observation/observe`。
每条绑定以 `event` 选择 `request_completed` 或 `websocket_response`,可分别配置请求范围;同一实例每种事件至多一条绑定,两类绑定使用相同 `order`。
事件选择只决定是否投递,不裁剪字段
@@ -50,7 +50,7 @@
## 扩展项简写
作者声明可以省略 `id` 和固定阶段;普通能力的版本默认为 `1`,`middleware` 显式声明版本 `4`,`upstream_adapter` 显式声明版本 `2`:
作者声明可以省略 `id` 和固定阶段;普通能力的版本默认为 `1`,`middleware` 显式声明版本 `4`,`scheduler` 和 `upstream_adapter` 显式声明版本 `2`:
```json
{
@@ -0,0 +1,60 @@
# 用量维护与计价
`call::services::usage` 提供请求明细、聚合统计和 Token 费用试算。插件通过 `call.host.service::<Operation>(input)` 调用,沿用公开服务的权限与调用期限
## 基础操作
| SDK 操作 | 服务名 | 输入与结果 |
| --- | --- | --- |
| `GetRecord` | `usage.records.get` | 按 `id` 读取原生请求明细,包含费用快照与 Provider 观测 |
| `ListRecords` | `usage.records.list` | 按时间、Provider、模型、账号、Key 筛选,返回明细与下一页游标 |
| `WriteRecords` | `usage.records.write` | 批量 `put`、`update`、`delete`,返回实际变更的 ID |
| `ListAggregates` | `usage.aggregates.list` | 分页读取调用方维护的聚合贡献,支持同样的时间与维度筛选 |
| `WriteAggregates` | `usage.aggregates.write` | 按聚合键批量覆盖或删除统计贡献 |
| `EvaluatePricing` | `pricing.evaluate` | 根据当前价目或指定价格试算,逐项返回 `estimated` 或 `not_estimable` |
原生总量与趋势通过既有统计查询读取。明细列表只列请求记录,聚合贡献参与统计但不生成请求明细
## 明细维护
- `put` 使用调用方提供的稳定 `id` 创建记录,已有 ID 则覆盖合同中的明细字段
- `update` 按 ID 替换全部 Token 字段或费用;未提供的部分保持原值,目标不存在返回 `not_found`
- `delete` 删除指定 ID;目标不存在时成功,且不加入 `changed_ids`
- 写入对象必须是已结束的请求,不能修改或删除仍在运行的请求
- `put` 和 `update` 清除原计价明细快照,后续读取使用写入的 Token 与费用
`UsageTokens` 中的缓存 Token 包含在输入 Token 中,推理 Token 包含在输出 Token 中。`null` 表示未知,`"0"` 表示已知为零;调用方负责从外部系统转换口径,宿主不核对各项业务关系
费用为 `unavailable` 或 `known { source, amount, currency }`。零费用属于已知费用;清空费用使用 `unavailable`。明细维护影响原生历史统计,不重新执行请求结算或修改 Key 预算账本
## 聚合维护
`UsageAggregateKey` 由 `bucket_start_ms`、`dimensions` 和调用方的 `key` 共同构成。相同完整键的 `put` 覆盖旧值,重复提交相同内容不会累计;改变时间、维度或 `key` 会形成另一个目标,需要移动时由调用方删除旧键
`dimensions` 支持 Provider、模型、账号、Key、操作和传输类型。`counts` 的键由 `UsageMetric` 定义,包括请求次数、各终态次数、已完成次数和各项 Token;`costs` 按币种保存十进制金额。计数与费用允许负值,供调用方表达统计修正
聚合按 `bucket_start_ms` 这个归属时间纳入 `[start_ms, end_ms)` 查询及趋势时间桶,宿主不自动摊分。导入日汇总时,调用方选择该日的归属时间;需要小时分布时,调用方自行提供小时数据
查询只按实际提供的维度匹配,缺少的维度不会自动补全。未提供某个维度的贡献参与总量,但不进入该维度的分组或筛选;分组与总量是否对齐由调用方提供的数据决定。请求 ID、响应 ID、状态码、请求结果及文字搜索等明细筛选不包含聚合贡献;聚合也不生成延迟样本、尝试次数或请求列表页数。原生统计的无符号计数与非负金额若因贡献合计超出表示范围,查询返回错误,由调用方修正输入
聚合贡献参与原生总量、趋势、维度统计和账号用量,统计重建保留这些数据。贡献由调用方删除,不随请求历史保留任务清理;明细沿用请求历史保留策略
## 费用试算与写回
`EvaluatePricing` 使用 Provider 的原生单请求 Token 计价公式,支持当前价目和调用方指定的价格覆盖。OpenAI 与 xAI 提供此能力;不支持的 Provider、缺少必要 Token、图像用量或无法确定的价目返回 `not_estimable`
试算输入需要明确输入、输出、缓存读取、缓存写入以及图像 Token;纯文本用量的图像 Token 填零。结果只包含 Token 费用,不自动处理图像、工具或其他附加费用,也不写入历史
调用方可查询明细、选择目标、逐条试算,再通过 `WriteRecords::Update` 写回费用。补全缺失时只处理 `unavailable`;全部覆盖时自行选择整个范围。日汇总若缺少逐请求分布,长上下文档位和逐请求舍入可能无法复现,计算方法由调用方决定
## 批量、错误与重试
- 每次写入或试算接受 1~500 项,编码后的输入不超过 2 MiB
- 查询每页默认 100 条,最多 200 条;明细页累计内容超过 2 MiB 时需减小页大小
- 时间为 UTC 毫秒,查询范围为左闭右开;游标不透明,仅用于相同筛选条件的后续页,分页不冻结跨调用快照
- `UsageCount`、`UsageDelta` 使用十进制字符串,避免 JavaScript 丢失整数精度;持久化数值受 PostgreSQL bigint 与费用精度范围限制
- 单次写入按输入顺序在一个事务中执行,任一项失败则整批回滚;跨批次进度由调用方保存
- 存储操作共享宿主观测连接与执行预算,超时、锁等待或取消会结束该次执行并释放资源
- 错误类型包括 `invalid_input`、`not_found`、`conflict`、`capacity_exceeded`、`unavailable`
响应丢失时,调用方按稳定 ID 或聚合键查询、重试覆盖。宿主不保存任务、来源登记、幂等回执、调整历史或撤销计划;调用方负责来源映射、业务去重、重试策略、并发写入协调和恢复。服务输入不能指定审计主体,宿主从当前调用身份生成普通变更审计
@@ -77,6 +77,8 @@ pub enum ModelRouteDecision {
#[serde(deny_unknown_fields)]
pub struct AccountScheduleCandidate {
pub account_id: String,
/// false 表示仅可等待,仍可被过滤,但不能作为本轮 Pick 的结果
pub available: bool,
pub weight: u16,
pub in_flight: u32,
pub maximum_concurrency: u32,
@@ -108,7 +110,13 @@ pub struct AccountScheduleRequest {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "decision", rename_all = "snake_case", deny_unknown_fields)]
pub enum AccountScheduleDecision {
Pick { account_id: String },
Pick {
account_id: String,
},
Delegate,
/// 只收窄本轮候选;参数与排除策略由插件拥有,宿主继续执行原生调度
DelegateFiltered {
exclude_account_ids: Vec<String>,
},
Reject,
}

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