mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-10 00:06:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
447869860a | ||
|
|
815c93a715 | ||
|
|
5843091969 | ||
|
|
1280c437b4 | ||
|
|
d965bc9050 | ||
|
|
bd9b017ed3 | ||
|
|
338cde1505 | ||
|
|
b25f96559a | ||
|
|
a8a96d8919 | ||
|
|
7b8339ede3 | ||
|
|
f439a30612 | ||
|
|
8572873c5f | ||
|
|
57c8ef358f | ||
|
|
70ab8ebc8d | ||
|
|
14ea5234ce | ||
|
|
f8480f3bdc |
Generated
+1
@@ -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);
|
||||
|
||||
@@ -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, §ion, "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, §ion, "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, §ion, "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");
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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, ¤t)
|
||||
.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, ¤t, &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))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user