feat(usage): expose bounded data maintenance services (#436)

Expose generic detail CRUD, aggregate contributions, and native pricing evaluation through Admin and the plugin SDK. Reuse existing request and statistics storage without schema changes; keep workflow state and recovery with callers.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-16 00:49:51 +08:00
committed by GitHub
co-authored by Zyy Codex
parent 8572873c5f
commit f439a30612
70 changed files with 3848 additions and 18 deletions
@@ -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);
+22
View File
@@ -39,6 +39,7 @@ pub use use_case::{
proxies::ProxiesService,
settings::SettingsService,
system::SystemService,
usage_data::UsageDataService,
};
use model::AdminError;
@@ -217,6 +218,7 @@ pub struct AdminServices {
client_distribution: Arc<dyn ClientDistributionService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
usage_data: Arc<dyn UsageDataService>,
system: Arc<dyn SystemService>,
credentials: Arc<CredentialsService>,
// Runtime 只持有 Weak,这个强引用保障账号端口与 Admin 服务同寿命
@@ -285,6 +287,11 @@ impl AdminServices {
self.client_distribution.as_ref()
}
#[must_use]
pub fn usage_data(&self) -> &dyn UsageDataService {
self.usage_data.as_ref()
}
#[must_use]
pub fn observability(&self) -> &dyn ObservabilityService {
self.observability.as_ref()
@@ -467,8 +474,10 @@ async fn initialize_inner(
registry.clone(),
pricing_source,
);
let usage_data = initialize_usage_data(store.usage_data(), store.settings(), registry.clone());
let mut public_services = public_service::Registry::new(service_middleware);
public_services.register_settings(&settings)?;
public_services.register_usage_data(&usage_data)?;
let services = AdminServices {
timezone,
public_services: Arc::new(public_services),
@@ -514,6 +523,7 @@ async fn initialize_inner(
timezone,
)),
settings,
usage_data,
system,
credentials,
_plugin_accounts: plugin_accounts,
@@ -644,3 +654,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,
))
}
@@ -24,6 +24,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)]
@@ -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 能力及账号类型,不执行网络或数据库查询
@@ -565,9 +565,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 +604,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 +616,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(),
}
}
@@ -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,188 @@
//! 公开原生用量读写与纯计价,插件拥有跨调用流程和恢复策略
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(&ListUsageRecords {
filter: request.filter.clone(),
cursor: request.cursor.clone(),
limit: 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)?;
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,192 @@
//! 用量输入沿用原生记录的数值、时间和资源边界
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(value: &ListUsageRecords) -> UsageDataResult<()> {
if value
.limit
.is_some_and(|v| v == 0 || v > USAGE_DATA_MAX_PAGE_SIZE)
|| value.filter.start_ms >= value.filter.end_ms
|| chrono::DateTime::from_timestamp_millis(value.filter.start_ms).is_none()
|| chrono::DateTime::from_timestamp_millis(value.filter.end_ms).is_none()
{
return Err(UsageDataError::invalid("查询时间范围或分页大小无效"));
}
for v in [
&value.filter.provider,
&value.filter.model,
&value.filter.account_ref,
&value.filter.client_key_ref,
]
.into_iter()
.flatten()
{
text(v, 128)?;
}
if value.cursor.as_ref().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(())
}
@@ -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(),
@@ -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");
}
@@ -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))),
@@ -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")
}
}
@@ -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 {
@@ -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)
}
}
@@ -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;
}
@@ -720,6 +720,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))
}
@@ -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
}
+3 -1
View File
@@ -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` 推进,
@@ -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 或聚合键查询、重试覆盖。宿主不保存任务、来源登记、幂等回执、调整历史或撤销计划;调用方负责来源映射、业务去重、重试策略、并发写入协调和恢复。服务输入不能指定审计主体,宿主从当前调用身份生成普通变更审计
@@ -1,6 +1,7 @@
//! 公开服务的类型化操作合同;名称用于分派,不承担授权
pub mod settings;
pub mod usage;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
@@ -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 operations;
pub use operations::*;
mod aggregate;
pub use aggregate::*;
@@ -0,0 +1,45 @@
//! 原生用量能力的类型化操作标识
use super::*;
use crate::call::services::Operation;
pub struct ListRecords;
impl Operation for ListRecords {
const NAME: &'static str = "usage.records.list";
type Input = ListUsageRecords;
type Output = UsageRecordPage;
}
pub struct GetRecord;
impl Operation for GetRecord {
const NAME: &'static str = "usage.records.get";
type Input = GetUsageRecord;
type Output = UsageRecord;
}
pub struct WriteRecords;
impl Operation for WriteRecords {
const NAME: &'static str = "usage.records.write";
type Input = WriteUsageRecords;
type Output = UsageWriteResult;
}
pub struct EvaluatePricing;
impl Operation for EvaluatePricing {
const NAME: &'static str = "pricing.evaluate";
type Input = EvaluateUsagePricing;
type Output = UsagePricingEvaluation;
}
pub struct ListAggregates;
impl Operation for ListAggregates {
const NAME: &'static str = "usage.aggregates.list";
type Input = ListUsageAggregates;
type Output = UsageAggregatePage;
}
pub struct WriteAggregates;
impl Operation for WriteAggregates {
const NAME: &'static str = "usage.aggregates.write";
type Input = WriteUsageAggregates;
type Output = UsageAggregateWriteResult;
}
@@ -0,0 +1,62 @@
//! 纯 Token 计价合同,费用是否采用与额外费用合并由调用方决定
use super::{UsageCost, UsageTokens};
use crate::call::services::settings::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)
}
}
@@ -23,6 +23,7 @@ fn diagnostics_do_not_render_peer_payloads_or_error_details() {
assert!(!format!("{frame:?} {fault:?}").contains(secret));
}
#[cfg(feature = "io")]
#[tokio::test]
async fn malformed_frame_metadata_retains_the_native_json_location() {
use std::error::Error as _;
@@ -172,6 +172,10 @@ async fn connect(
)?);
let plugins = Arc::new(postgres::PgPluginStore::new(pool.clone()));
let usage_data = Arc::new(postgres::PgUsageDataStore::new(
pool.clone(),
observability_query_budget.clone(),
));
let admin_ports = AdminStorePorts::new(
AdminAccountStorePorts::new(
Arc::new(postgres::PgAdminAccountStore::new(
@@ -212,6 +216,7 @@ async fn connect(
plugins.clone(),
plugins.clone(),
plugins,
usage_data,
);
let execution_repository = Arc::new(postgres::PgExecutionStore::new(pool.clone()));
@@ -20,6 +20,7 @@ mod query_pattern;
mod retention;
mod runtime_settings;
mod snapshot;
mod usage_data;
mod usage_facts;
pub use account_groups::*;
@@ -44,6 +45,7 @@ use query_pattern::literal_prefix_pattern;
pub use retention::*;
pub use runtime_settings::*;
pub use snapshot::*;
pub use usage_data::PgUsageDataStore;
pub(crate) use usage_facts::{
completed_usage_fact_predicate, push_completed_usage_fact_filter,
push_unrecovered_request_filter,
@@ -26,6 +26,58 @@ impl AnalyticsSnapshot<'_> {
range: ObservabilityRange,
account_ids: &[String],
by_model: bool,
) -> StoreResult<AccountStates> {
let native = self
.native_account_states(range, account_ids, by_model)
.await?;
let mut states = BTreeMap::new();
for (key, value) in native {
states.insert(
(
range.start,
serde_json::to_string(&key)
.map_err(|_| invalid("encode aggregate account key"))?,
),
value,
);
}
let filter = UsageRecordFilter::default();
self.merge_contributions(
super::contributions::ContributionQuery {
range,
filter: &filter,
dimension: None,
granularity: None,
timezone: Default::default(),
account_models: Some(by_model),
selection: super::query::ReadSelection {
distributions: false,
costs: true,
keys: Some(account_ids),
counters: None,
distribution_names: None,
},
},
&mut states,
)
.await?;
states
.into_iter()
.map(|((_, key), value)| {
Ok((
serde_json::from_str(&key)
.map_err(|_| invalid("decode aggregate account key"))?,
value,
))
})
.collect()
}
async fn native_account_states(
&mut self,
range: ObservabilityRange,
account_ids: &[String],
by_model: bool,
) -> StoreResult<AccountStates> {
if account_ids.is_empty() {
return Ok(BTreeMap::new());
@@ -97,8 +97,14 @@ pub(super) async fn rebuild_bucket_batch(
"request_analytics_1h"
};
let start = if phase == 0 { bucket } else { hour };
// 调用方贡献独立保存在小时表,完整小时查询不额外依赖十五分钟表
let retained = if phase == 3 {
" and collection<>'usage_contribution'"
} else {
""
};
let statement = format!(
"delete from {table} where ctid in (select ctid from {table} where bucket_start=$1 limit $2)"
"delete from {table} where ctid in (select ctid from {table} where bucket_start=$1{retained} limit $2)"
);
let removed = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(start)
@@ -0,0 +1,223 @@
//! 将调用方维护的聚合贡献合入原生统计,保留查询预算与数值边界
use super::super::super::usage_data::AGGREGATE_COLLECTION;
use super::query::ReadSelection;
use super::*;
use futures::TryStreamExt;
pub(super) struct ContributionQuery<'a> {
pub range: ObservabilityRange,
pub filter: &'a UsageRecordFilter,
pub dimension: Option<DiagnosticDimension>,
pub granularity: Option<ObservationGranularity>,
pub timezone: gateway_core::time::DeploymentTimeZone,
pub account_models: Option<bool>,
pub selection: ReadSelection<'a>,
}
impl AnalyticsSnapshot<'_> {
pub(super) async fn merge_contributions(
&mut self,
input: ContributionQuery<'_>,
states: &mut BTreeMap<(DateTime<Utc>, String), MetricState>,
) -> StoreResult<()> {
let ContributionQuery {
range,
filter,
dimension,
granularity,
timezone,
account_models,
selection,
} = input;
// 聚合未提供请求级身份或状态维度,不能用猜测的明细值匹配这些筛选
if filter.request_id.is_some()
|| filter.response_id.is_some()
|| filter.upstream_request_id.is_some()
|| filter.attempt_index.is_some()
|| filter.status_code.is_some()
|| filter.outcome.is_some()
|| filter.search.is_some()
|| matches!(
dimension,
Some(DiagnosticDimension::Failure | DiagnosticDimension::Status)
)
{
return Ok(());
}
let name = if let Some(by_model) = account_models {
if by_model {
"jsonb_build_array(dimensions[3],dimensions[2])::text"
} else {
"jsonb_build_array(dimensions[3],'')::text"
}
} else {
match dimension {
None => "''::text",
Some(DiagnosticDimension::Provider) => {
"coalesce(nullif(dimensions[1],''),'unrouted')"
}
Some(DiagnosticDimension::Model) => "dimensions[2]",
Some(DiagnosticDimension::Account) => {
"coalesce(nullif(dimensions[3],''),'unrouted')"
}
Some(DiagnosticDimension::ApiKey) => "dimensions[4]",
Some(DiagnosticDimension::Transport) => {
"coalesce(nullif(dimensions[6],''),'unknown')"
}
_ => return Ok(()),
}
};
let boundaries = if granularity == Some(ObservationGranularity::Day) {
Some(calendar_buckets(range, timezone)?)
} else {
None
};
let mut q = QueryBuilder::<Postgres>::new("with matching as materialized (select ");
aggregate::push_output_bucket(
&mut q,
"bucket_start",
range,
granularity,
boundaries.as_deref(),
);
q.push(format!(" as output_bucket,{name} as name,counts,costs from request_analytics_1h where collection=")).push_bind(AGGREGATE_COLLECTION).push(" and bucket_start >= ").push_bind(range.start).push(" and bucket_start < ").push_bind(range.end);
for (index, value) in [
(1, &filter.provider_kind),
(2, &filter.model),
(3, &filter.provider_account_ref),
(4, &filter.client_api_key_ref),
(5, &filter.operation),
(6, &filter.transport),
] {
if let Some(value) = value {
q.push(format!(" and dimensions[{index}]="))
.push_bind(value);
}
}
if let Some(keys) = selection.keys {
let key = if account_models.is_some() {
"dimensions[3]"
} else {
name
};
q.push(format!(" and ({key})=any("))
.push_bind(keys)
.push(")");
}
let dimension_index = match dimension {
Some(DiagnosticDimension::Provider) => Some(1),
Some(DiagnosticDimension::Model) => Some(2),
Some(DiagnosticDimension::Account) => Some(3),
Some(DiagnosticDimension::ApiKey) => Some(4),
Some(DiagnosticDimension::Transport) => Some(6),
_ => None,
};
if let Some(index) = dimension_index {
// 缺少归属维度的贡献只进入总量,不分配给虚构的未知分组
q.push(format!(" and dimensions[{index}]<>''"));
}
q.push("), entries as (select output_bucket,name,'count' as kind,c.key,c.value::numeric as value from matching cross join lateral jsonb_each_text(counts) c");
if let Some(counters) = selection.counters {
q.push(" where c.key=any(").push_bind(counters).push(")");
}
if selection.costs {
q.push(" union all select output_bucket,name,'cost',c.key,c.value::numeric from matching cross join lateral jsonb_each_text(costs) c");
}
q.push("), totals as (select output_bucket,name,kind,key,sum(value)::text as value from entries group by output_bucket,name,kind,key) select output_bucket,name,jsonb_object_agg(key,value) filter(where kind='count') as counts,jsonb_object_agg(key,value) filter(where kind='cost') as costs from totals group by output_bucket,name");
let mut rows = q.build().fetch(&mut *self.tx);
while let Some(row) = rows
.try_next()
.await
.map_err(|e| postgres_unavailable("read maintained aggregate contributions", e))?
{
let mut name: String = get(&row, "name")?;
if account_models.is_some() {
let key: (String, String) = serde_json::from_str(&name)
.map_err(|_| invalid("invalid aggregate account key"))?;
name = serde_json::to_string(&key)
.map_err(|_| invalid("encode aggregate account key"))?;
}
let key = (get(&row, "output_bucket")?, name);
let state = states.entry(key).or_default();
let counts =
get::<Option<sqlx::types::Json<BTreeMap<String, String>>>>(&row, "counts")?
.map(|v| v.0)
.unwrap_or_default();
for (name, value) in &counts {
let delta = value
.parse::<i128>()
.map_err(|_| invalid("aggregate counter exceeds arithmetic range"))?;
apply_count(state, name, delta)?;
if metrics::TOKEN_METRICS.contains(&name.as_str()) {
// 只标记指标已知,不虚构请求明细、延迟样本或费用覆盖率
state.increment(&format!("{name}_known_count"), 1)?;
}
}
if selection.counters.is_none() {
let total = counts
.get("total_tokens")
.map(|v| v.parse::<i128>())
.transpose()
.map_err(|_| invalid("aggregate token overflow"))?;
let token_delta = match total {
Some(total) => total,
None => {
let input = counts
.get("input_tokens")
.map(|v| v.parse::<i128>())
.transpose()
.map_err(|_| invalid("aggregate input overflow"))?
.unwrap_or_default();
let output = counts
.get("output_tokens")
.map(|v| v.parse::<i128>())
.transpose()
.map_err(|_| invalid("aggregate output overflow"))?
.unwrap_or_default();
input
.checked_add(output)
.ok_or_else(|| invalid("aggregate token overflow"))?
}
};
apply_count(state, "window_total_tokens", token_delta)?;
}
let costs = get::<Option<sqlx::types::Json<BTreeMap<String, String>>>>(&row, "costs")?
.map(|v| v.0)
.unwrap_or_default();
for (currency, value) in costs {
let current = state
.costs
.get(&currency)
.map(String::as_str)
.unwrap_or("0");
let current = current
.parse::<gateway_core::metering::Decimal>()
.map_err(|_| invalid("native cost exceeds decimal range"))?;
let magnitude = value
.strip_prefix('-')
.unwrap_or(&value)
.parse::<gateway_core::metering::Decimal>()
.map_err(|_| invalid("aggregate cost exceeds decimal range"))?;
let scaled = if value.starts_with('-') {
current.scaled().checked_sub(magnitude.scaled())
} else {
current.scaled().checked_add(magnitude.scaled())
}
.ok_or_else(|| invalid("aggregate cost is outside native result range"))?;
let amount = gateway_core::metering::Decimal::from_scaled(scaled)
.map_err(|_| invalid("aggregate cost exceeds native result range"))?;
state.costs.insert(currency, amount.canonical());
}
}
Ok(())
}
}
fn apply_count(state: &mut MetricState, name: &str, delta: i128) -> StoreResult<()> {
let total = i128::from(state.count(name))
.checked_add(delta)
.and_then(|v| u64::try_from(v).ok())
.ok_or_else(|| invalid("aggregate counter is outside native result range"))?;
state.counts.insert(name.into(), total);
Ok(())
}
@@ -3,6 +3,7 @@
pub(crate) mod accounts;
mod aggregate;
mod buckets;
mod contributions;
mod counts;
mod forecast;
mod histogram;
@@ -9,12 +9,13 @@ pub(crate) struct AnalyticsSnapshot<'a> {
approximate: bool,
}
struct ReadSelection<'a> {
distributions: bool,
costs: bool,
keys: Option<&'a [String]>,
counters: Option<&'a [&'a str]>,
distribution_names: Option<&'a [&'a str]>,
#[derive(Clone, Copy)]
pub(super) struct ReadSelection<'a> {
pub(super) distributions: bool,
pub(super) costs: bool,
pub(super) keys: Option<&'a [String]>,
pub(super) counters: Option<&'a [&'a str]>,
pub(super) distribution_names: Option<&'a [&'a str]>,
}
impl<'a> AnalyticsSnapshot<'a> {
@@ -107,7 +108,7 @@ impl<'a> AnalyticsSnapshot<'a> {
pub async fn totals(&mut self) -> StoreResult<DashboardTotals> {
self.require_visible(None)?;
let bounds=sqlx::query("select (select started_at from request_analytics_facts order by started_at,id limit 1) as first_started_at,(select started_at from request_analytics_facts order by started_at desc,id desc limit 1) as last_started_at")
let bounds=sqlx::query("select least((select started_at from request_analytics_facts order by started_at,id limit 1),(select bucket_start from request_analytics_1h where collection='usage_contribution' order by bucket_start limit 1)) as first_started_at,greatest((select started_at from request_analytics_facts order by started_at desc,id desc limit 1),(select bucket_start from request_analytics_1h where collection='usage_contribution' order by bucket_start desc limit 1)) as last_started_at")
.fetch_one(&mut *self.tx).await.map_err(|error|postgres_unavailable("read retained analytics bounds",error))?;
let start =
get::<Option<DateTime<Utc>>>(&bounds, "first_started_at")?.unwrap_or_else(Utc::now);
@@ -363,7 +364,10 @@ impl<'a> AnalyticsSnapshot<'a> {
first_token_p95_ms: state
.percentile("first_token_ms", 0.95)
.map(|value| value.round() as u64),
non_completion_count: metrics.cancelled_count + metrics.incomplete_count,
non_completion_count: metrics
.cancelled_count
.checked_add(metrics.incomplete_count)
.ok_or_else(|| invalid("non-completion count overflow"))?,
retry_count: state.count("retry_count"),
retried_request_count: state.count("retried_request_count"),
cost_coverage: state.coverage(),
@@ -477,6 +481,21 @@ impl<'a> AnalyticsSnapshot<'a> {
}
query.push(format!(" group by {column}"));
}
if filter.outcome.is_none() {
// 贡献也在数据库内参与 TopN,避免有补录时退回加载全部账号或 Key
let index = if dimension == DiagnosticDimension::Account {
3
} else {
4
};
query.push(format!(" union all select dimensions[{index}] as key,coalesce((counts->>'request_count')::bigint,0) as request_count from request_analytics_1h where collection='usage_contribution' and dimensions[{index}]<>'' and bucket_start >= "))
.push_bind(range.start).push(" and bucket_start < ").push_bind(range.end);
if let Some(reference) = reference {
query
.push(format!(" and dimensions[{index}]="))
.push_bind(reference);
}
}
query.push("), grouped as (select key,sum(request_count)::bigint as request_count from contributions group by key) select key,request_count,sum(request_count) over()::bigint as total_request_count from grouped order by request_count desc,");
if dimension == DiagnosticDimension::Account {
query.push("coalesce(nullif(key,''),'unrouted')");
@@ -600,6 +619,34 @@ impl<'a> AnalyticsSnapshot<'a> {
granularity: Option<ObservationGranularity>,
timezone: gateway_core::time::DeploymentTimeZone,
selection: ReadSelection<'_>,
) -> StoreResult<BTreeMap<(DateTime<Utc>, String), MetricState>> {
let mut states = self
.load_native_selected(range, filter, dimension, granularity, timezone, selection)
.await?;
self.merge_contributions(
super::contributions::ContributionQuery {
range,
filter,
dimension,
granularity,
timezone,
account_models: None,
selection,
},
&mut states,
)
.await?;
Ok(states)
}
async fn load_native_selected(
&mut self,
range: ObservabilityRange,
filter: &UsageRecordFilter,
dimension: Option<DiagnosticDimension>,
granularity: Option<ObservationGranularity>,
timezone: gateway_core::time::DeploymentTimeZone,
selection: ReadSelection<'_>,
) -> StoreResult<BTreeMap<(DateTime<Utc>, String), MetricState>> {
let ReadSelection {
distributions: read_distributions,
@@ -57,14 +57,18 @@ pub(super) async fn reset_obsolete_batch(
_ => return Err(invalid("invalid analytics reset phase")),
};
let mut query = QueryBuilder::<Postgres>::new(format!(
"with batch as materialized (select ctid,{keys} from {table}"
"with batch as materialized (select ctid,{keys} from {table} where true"
));
// 手工维护的统计贡献是输入数据,不属于可丢弃的请求投影
if phase == 3 {
query.push(" and collection<>'usage_contribution'");
}
let cursor = get::<Option<sqlx::types::Json<ResetCursor>>>(&state, "reset_cursor")?;
// LIMIT 只限制删除数量,主键下界才避免后续批次反复扫描已经清空的堆页
match cursor.map(|cursor| cursor.0) {
None => {}
Some(ResetCursor::Fact { id }) if phase == 1 => {
query.push(" where id > ").push_bind(id);
query.push(" and id > ").push_bind(id);
}
Some(ResetCursor::Rollup {
bucket_start,
@@ -73,7 +77,7 @@ pub(super) async fn reset_obsolete_batch(
outcome,
}) if phase == 2 || phase == 3 => {
query
.push(" where (bucket_start,collection,dimensions,outcome) > (")
.push(" and (bucket_start,collection,dimensions,outcome) > (")
.push_bind(bucket_start)
.push(",")
.push_bind(collection)
@@ -84,7 +88,7 @@ pub(super) async fn reset_obsolete_batch(
.push(")");
}
Some(ResetCursor::Bucket { bucket_start }) if phase == 4 => {
query.push(" where bucket_start > ").push_bind(bucket_start);
query.push(" and bucket_start > ").push_bind(bucket_start);
}
Some(_) => return Err(invalid("analytics reset cursor does not match its phase")),
}
@@ -0,0 +1,129 @@
//! 在既有聚合表中维护调用方贡献,主键覆盖不追加去重账本
use super::*;
use sqlx::QueryBuilder;
use std::collections::BTreeMap;
pub(crate) const AGGREGATE_COLLECTION: &str = "usage_contribution";
fn dimensions(key: &UsageAggregateKey) -> Vec<String> {
let d = &key.dimensions;
[
&d.provider,
&d.model,
&d.account_ref,
&d.client_key_ref,
&d.operation,
&d.transport,
]
.into_iter()
.map(|v| v.clone().unwrap_or_default())
.chain(std::iter::once(key.key.clone()))
.collect()
}
pub(super) async fn write(
tx: &mut Transaction<'_, Postgres>,
change: UsageAggregateChange,
) -> StoreResult<UsageAggregateKey> {
match change {
UsageAggregateChange::Delete { identity } => {
sqlx::query("delete from request_analytics_1h where collection=$1 and bucket_start=$2 and dimensions=$3 and outcome=''").bind(AGGREGATE_COLLECTION).bind(at(identity.bucket_start_ms)?).bind(dimensions(&identity)).execute(&mut **tx).await.map_err(db_error)?;
Ok(identity)
}
UsageAggregateChange::Put { aggregate } => {
let counts: BTreeMap<_, _> = aggregate
.counts
.iter()
.map(|(metric, value)| (*metric, value.get()))
.collect();
sqlx::query("insert into request_analytics_1h(bucket_start,collection,dimensions,outcome,schema_version,histogram_version,counts,costs,distributions) values($1,$2,$3,'',0,0,$4,$5,'{}') on conflict(bucket_start,collection,dimensions,outcome) do update set counts=excluded.counts,costs=excluded.costs")
.bind(at(aggregate.identity.bucket_start_ms)?).bind(AGGREGATE_COLLECTION).bind(dimensions(&aggregate.identity)).bind(sqlx::types::Json(counts)).bind(sqlx::types::Json(aggregate.costs)).execute(&mut **tx).await.map_err(db_error)?;
Ok(aggregate.identity)
}
}
}
pub(super) async fn list(
tx: &mut Transaction<'_, Postgres>,
request: ListUsageAggregates,
) -> StoreResult<UsageAggregatePage> {
let limit = request.limit.unwrap_or(100).min(USAGE_DATA_MAX_PAGE_SIZE);
let mut q = QueryBuilder::new(
"select bucket_start,dimensions,counts,costs from request_analytics_1h where collection=",
);
q.push_bind(AGGREGATE_COLLECTION)
.push(" and bucket_start >= ")
.push_bind(at(request.filter.start_ms)?)
.push(" and bucket_start < ")
.push_bind(at(request.filter.end_ms)?);
for (index, value) in [
(1, request.filter.provider),
(2, request.filter.model),
(3, request.filter.account_ref),
(4, request.filter.client_key_ref),
] {
if let Some(value) = value {
q.push(format!(" and dimensions[{index}]="))
.push_bind(value);
}
}
if let Some(cursor) = request.cursor {
let identity: UsageAggregateKey =
serde_json::from_str(&cursor).map_err(|_| invalid("invalid aggregate cursor"))?;
q.push(" and (bucket_start,dimensions) > (")
.push_bind(at(identity.bucket_start_ms)?)
.push(",")
.push_bind(dimensions(&identity))
.push(")");
}
q.push(" order by bucket_start,dimensions limit ")
.push_bind(i64::from(limit) + 1);
let mut rows = q.build().fetch_all(&mut **tx).await.map_err(db_error)?;
let more = rows.len() > usize::from(limit);
rows.truncate(usize::from(limit));
let mut items = Vec::with_capacity(rows.len());
for row in rows {
let d: Vec<String> = row.try_get("dimensions").map_err(db_error)?;
if d.len() != 7 {
return Err(invalid("invalid aggregate dimensions"));
}
let optional = |i: usize| (!d[i].is_empty()).then(|| d[i].clone());
let counts: sqlx::types::Json<BTreeMap<UsageMetric, i64>> =
row.try_get("counts").map_err(db_error)?;
let costs: sqlx::types::Json<BTreeMap<String, String>> =
row.try_get("costs").map_err(db_error)?;
items.push(UsageAggregate {
identity: UsageAggregateKey {
bucket_start_ms: row
.try_get::<DateTime<Utc>, _>("bucket_start")
.map_err(db_error)?
.timestamp_millis(),
dimensions: UsageDimensions {
provider: optional(0),
model: optional(1),
account_ref: optional(2),
client_key_ref: optional(3),
operation: optional(4),
transport: optional(5),
},
key: d[6].clone(),
},
counts: counts
.0
.into_iter()
.map(|(k, v)| (k, UsageDelta(v)))
.collect(),
costs: costs.0,
});
}
let next_cursor = if more {
items
.last()
.map(|item| {
serde_json::to_string(&item.identity)
.map_err(|_| invalid("encode aggregate cursor"))
})
.transpose()?
} else {
None
};
Ok(UsageAggregatePage { items, next_cursor })
}
@@ -0,0 +1,173 @@
//! 复用请求明细和聚合存储的有界维护入口,不管理插件工作流
mod aggregates;
mod records;
pub(crate) use aggregates::AGGREGATE_COLLECTION;
use super::{ObservabilityConnection, ObservabilityQueryBudget};
use crate::{StoreError, StoreResult, postgres_unavailable};
use chrono::{DateTime, Utc};
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
use sqlx::{Connection as _, PgPool, Postgres, Row, Transaction};
#[derive(Clone)]
pub struct PgUsageDataStore {
pool: PgPool,
budget: ObservabilityQueryBudget,
}
impl PgUsageDataStore {
#[must_use]
pub const fn new(pool: PgPool, budget: ObservabilityQueryBudget) -> Self {
Self { pool, budget }
}
}
#[async_trait::async_trait]
impl UsageDataStore for PgUsageDataStore {
async fn get_record(&self, request: GetUsageRecord) -> UsageDataResult<UsageRecord> {
self.budget
.run(&self.pool, "usage record", |mut connection| async move {
let mut tx = begin(&mut connection).await?;
records::get(&mut tx, &request.id).await
})
.await
.map_err(public_error)
}
async fn list_records(&self, request: ListUsageRecords) -> UsageDataResult<UsageRecordPage> {
self.budget
.run(&self.pool, "usage records", |mut connection| async move {
let mut tx = begin(&mut connection).await?;
records::list(&mut tx, request).await
})
.await
.map_err(public_error)
}
async fn write_records(&self, command: StoredUsageWrite) -> UsageDataResult<UsageWriteResult> {
self.budget
.run(
&self.pool,
"write usage records",
|mut connection| async move {
let mut tx = begin(&mut connection).await?;
let mut changed_ids = Vec::with_capacity(command.changes.len());
for (change, event) in command.changes {
if records::write(&mut tx, &change).await? {
changed_ids.push(change.id().to_owned());
}
audit(&mut tx, event).await?;
}
tx.commit().await.map_err(db_error)?;
Ok(UsageWriteResult { changed_ids })
},
)
.await
.map_err(public_error)
}
async fn list_aggregates(
&self,
request: ListUsageAggregates,
) -> UsageDataResult<UsageAggregatePage> {
self.budget
.run(
&self.pool,
"usage aggregates",
|mut connection| async move {
let mut tx = begin(&mut connection).await?;
aggregates::list(&mut tx, request).await
},
)
.await
.map_err(public_error)
}
async fn write_aggregates(
&self,
command: StoredAggregateWrite,
) -> UsageDataResult<UsageAggregateWriteResult> {
self.budget.run(&self.pool, "write usage aggregates", |mut connection| async move {
let mut tx = begin(&mut connection).await?;
let mut changed=Vec::with_capacity(command.changes.len());
for (change,event) in command.changes {
changed.push(aggregates::write(&mut tx,change).await?);
audit(&mut tx,event).await?;
}
// 沿用现有统计发布代次,让读缓存能观察到直接维护的聚合值
sqlx::query("update request_analytics_build_state set generation=generation+1,updated_at=clock_timestamp() where singleton").execute(&mut *tx).await.map_err(db_error)?;
tx.commit().await.map_err(db_error)?;
Ok(UsageAggregateWriteResult { changed })
}).await.map_err(public_error)
}
}
async fn begin(connection: &mut ObservabilityConnection) -> StoreResult<Transaction<'_, Postgres>> {
let timeout = connection.query_timeout();
let mut tx = connection.begin().await.map_err(db_error)?;
sqlx::query(
"select set_config('statement_timeout',$1,true),set_config('lock_timeout',$1,true)",
)
.bind(format!("{}ms", timeout.as_millis()))
.execute(&mut *tx)
.await
.map_err(db_error)?;
Ok(tx)
}
fn at(ms: i64) -> StoreResult<DateTime<Utc>> {
DateTime::from_timestamp_millis(ms).ok_or_else(|| invalid("invalid usage time"))
}
fn invalid(message: &str) -> StoreError {
StoreError::InvalidData {
entity: "usage data",
message: message.into(),
source: None,
}
}
fn db_error(error: sqlx::Error) -> StoreError {
if error
.as_database_error()
.and_then(|e| e.code())
.is_some_and(|code| code.starts_with("22") || code.starts_with("23"))
{
return invalid("usage value violates storage constraints").with_source(error);
}
postgres_unavailable("usage storage operation", error)
}
fn public_error(error: StoreError) -> UsageDataError {
let kind = match error {
StoreError::InvalidData { .. } => UsageErrorKind::InvalidInput,
StoreError::NotFound { .. } => UsageErrorKind::NotFound,
StoreError::Conflict { .. } => UsageErrorKind::Conflict,
StoreError::Unavailable { .. } => UsageErrorKind::Unavailable,
};
tracing::debug!(%error,"usage operation failed");
UsageDataError::new(
kind,
match kind {
UsageErrorKind::InvalidInput => "数据不符合存储格式或数值范围",
UsageErrorKind::NotFound => "记录不存在",
UsageErrorKind::Conflict => "请求仍在运行,无法维护",
_ => "存储操作失败或超出资源预算",
},
)
}
async fn audit(
tx: &mut Transaction<'_, Postgres>,
event: gateway_admin::model::auth::AdminAuditEvent,
) -> StoreResult<()> {
super::append_admin_audit_event_in_transaction(
tx,
super::AdminAuditEvent {
id: event.id,
actor_kind: event.actor_kind.into(),
actor_admin_user_id: event.actor_admin_user_id,
actor_ref: event.actor_ref,
admin_request_id: event.request_id,
action: event.action,
entity_kind: event.entity_kind,
entity_ref: event.entity_ref,
config_revision: None,
changed_fields: event.changed_fields,
created_at: event.occurred_at,
},
None,
)
.await
}
@@ -0,0 +1,383 @@
//! 按原生主键维护已结束的请求明细,复用统计失效触发器
use super::*;
use futures::TryStreamExt;
use serde::{Deserialize, Serialize};
use sqlx::{QueryBuilder, postgres::PgRow};
const SELECT: &str = "select id,client_api_key_ref,provider_kind,provider_account_ref,requested_model_id,upstream_model_id,request_observation_json #>> '{request,protocol}' as protocol,request_observation_json #>> '{request,endpoint}' as endpoint,operation,client_transport,service_tier,started_at,completed_at,outcome,input_tokens,output_tokens,cached_input_tokens,cache_write_input_tokens,reasoning_output_tokens,image_input_tokens,image_output_tokens,total_tokens,cost_source,cost_currency,cost_amount::text as cost_amount,billing_snapshot_json,provider_observation_json from model_requests";
#[derive(Serialize, Deserialize)]
struct Cursor {
started_at: DateTime<Utc>,
id: String,
}
pub(super) async fn get(tx: &mut Transaction<'_, Postgres>, id: &str) -> StoreResult<UsageRecord> {
let row = sqlx::query(sqlx::AssertSqlSafe(format!("{SELECT} where id=$1")))
.bind(id)
.fetch_optional(&mut **tx)
.await
.map_err(db_error)?
.ok_or_else(|| StoreError::NotFound {
entity: "usage record",
id: id.into(),
source: None,
})?;
decode(row)
}
pub(super) async fn list(
tx: &mut Transaction<'_, Postgres>,
request: ListUsageRecords,
) -> StoreResult<UsageRecordPage> {
let limit = request.limit.unwrap_or(100).min(USAGE_DATA_MAX_PAGE_SIZE);
let mut q = QueryBuilder::new(SELECT);
q.push(" where started_at >= ")
.push_bind(at(request.filter.start_ms)?)
.push(" and started_at < ")
.push_bind(at(request.filter.end_ms)?);
for (column, value) in [
("provider_kind", request.filter.provider),
(
"coalesce(upstream_model_id,requested_model_id)",
request.filter.model,
),
("provider_account_ref", request.filter.account_ref),
("client_api_key_ref", request.filter.client_key_ref),
] {
if let Some(value) = value {
q.push(format!(" and {column}=")).push_bind(value);
}
}
if let Some(cursor) = request.cursor {
let cursor: Cursor =
serde_json::from_str(&cursor).map_err(|_| invalid("invalid usage cursor"))?;
q.push(" and (started_at,id) > (")
.push_bind(cursor.started_at)
.push(",")
.push_bind(cursor.id)
.push(")");
}
q.push(" order by started_at,id limit ")
.push_bind(i64::from(limit) + 1);
let mut rows = q.build().fetch(&mut **tx);
let mut items = Vec::new();
let mut last = None;
let mut bytes = 0usize;
let mut next_cursor = None;
while let Some(row) = rows.try_next().await.map_err(db_error)? {
if items.len() == usize::from(limit) {
next_cursor = last
.map(|v| serde_json::to_string(&v).map_err(|_| invalid("encode usage cursor")))
.transpose()?;
break;
}
last = Some(Cursor {
started_at: row.try_get("started_at").map_err(db_error)?,
id: row.try_get("id").map_err(db_error)?,
});
let record = decode(row)?;
bytes = bytes
.checked_add(
serde_json::to_vec(&record)
.map_err(|_| invalid("encode usage record"))?
.len(),
)
.ok_or_else(|| invalid("usage page byte overflow"))?;
// 原生观测 JSON 可能较大,流式读取并限制页面,避免先物化两百条大记录
if bytes > USAGE_DATA_MAX_REQUEST_BYTES {
return Err(invalid("usage page exceeds 2 MiB; reduce page size"));
}
items.push(record);
}
Ok(UsageRecordPage { items, next_cursor })
}
fn decode(row: PgRow) -> StoreResult<UsageRecord> {
let token = |name: &str| -> StoreResult<Option<UsageCount>> {
row.try_get::<Option<i64>, _>(name)
.map_err(db_error)?
.map(|v| {
u64::try_from(v)
.map(UsageCount)
.map_err(|_| invalid("negative native token count"))
})
.transpose()
};
let outcome = match row.try_get::<&str, _>("outcome").map_err(db_error)? {
"running" => UsageOutcome::Running,
"succeeded" => UsageOutcome::Succeeded,
"failed" => UsageOutcome::Failed,
"cancelled" => UsageOutcome::Cancelled,
"incomplete" => UsageOutcome::Incomplete,
_ => return Err(invalid("unknown native outcome")),
};
let source: String = row.try_get("cost_source").map_err(db_error)?;
let cost = match (
row.try_get::<Option<String>, _>("cost_amount")
.map_err(db_error)?,
row.try_get::<Option<String>, _>("cost_currency")
.map_err(db_error)?,
) {
(Some(amount), Some(currency)) => UsageCost::Known {
source: if source == "provider_reported" {
UsageCostSource::ProviderReported
} else {
UsageCostSource::Calculated
},
amount,
currency,
},
_ => UsageCost::Unavailable,
};
Ok(UsageRecord {
id: row.try_get("id").map_err(db_error)?,
data: UsageRecordData {
client_key_ref: row.try_get("client_api_key_ref").map_err(db_error)?,
provider: row.try_get("provider_kind").map_err(db_error)?,
account_ref: row.try_get("provider_account_ref").map_err(db_error)?,
requested_model: row.try_get("requested_model_id").map_err(db_error)?,
model: row.try_get("upstream_model_id").map_err(db_error)?,
protocol: row.try_get("protocol").map_err(db_error)?,
operation: row.try_get("operation").map_err(db_error)?,
endpoint: row.try_get("endpoint").map_err(db_error)?,
client_transport: row.try_get("client_transport").map_err(db_error)?,
service_tier: row.try_get("service_tier").map_err(db_error)?,
started_at_ms: row
.try_get::<DateTime<Utc>, _>("started_at")
.map_err(db_error)?
.timestamp_millis(),
completed_at_ms: row
.try_get::<Option<DateTime<Utc>>, _>("completed_at")
.map_err(db_error)?
.map(|v| v.timestamp_millis()),
outcome,
tokens: UsageTokens {
input_tokens: token("input_tokens")?,
output_tokens: token("output_tokens")?,
cached_input_tokens: token("cached_input_tokens")?,
cache_write_input_tokens: token("cache_write_input_tokens")?,
reasoning_output_tokens: token("reasoning_output_tokens")?,
image_input_tokens: token("image_input_tokens")?,
image_output_tokens: token("image_output_tokens")?,
total_tokens: token("total_tokens")?,
},
cost,
},
billing_snapshot: row.try_get("billing_snapshot_json").map_err(db_error)?,
provider_observation: row.try_get("provider_observation_json").map_err(db_error)?,
})
}
fn cost_parts(cost: &UsageCost) -> (&str, Option<&str>, Option<&str>) {
match cost {
UsageCost::Unavailable => ("unavailable", None, None),
UsageCost::Known {
source,
amount,
currency,
} => (
match source {
UsageCostSource::ProviderReported => "provider_reported",
UsageCostSource::Calculated => "calculated",
},
Some(amount),
Some(currency),
),
}
}
fn token_values(t: &UsageTokens) -> StoreResult<Vec<Option<i64>>> {
[
t.input_tokens,
t.output_tokens,
t.cached_input_tokens,
t.cache_write_input_tokens,
t.reasoning_output_tokens,
t.image_input_tokens,
t.image_output_tokens,
t.total_tokens,
]
.into_iter()
.map(|v| {
v.map(|v| i64::try_from(v.get()).map_err(|_| invalid("native token overflow")))
.transpose()
})
.collect()
}
pub(super) async fn write(
tx: &mut Transaction<'_, Postgres>,
change: &UsageRecordChange,
) -> StoreResult<bool> {
// 请求执行持有自己的状态;维护终态数据不参与运行请求的结算或租约
let running: Option<String> =
sqlx::query_scalar("select outcome from model_requests where id=$1 for update")
.bind(change.id())
.fetch_optional(&mut **tx)
.await
.map_err(db_error)?;
if running.as_deref() == Some("running") {
return Err(StoreError::Conflict {
entity: "usage record",
id: change.id().into(),
kind: crate::ConflictKind::InvalidTransition,
source: None,
});
}
match change {
UsageRecordChange::Delete { id } => Ok(sqlx::query(
"delete from model_requests where id=$1 and outcome<>'running'",
)
.bind(id)
.execute(&mut **tx)
.await
.map_err(db_error)?
.rows_affected()
> 0),
UsageRecordChange::Put { id, data } => {
let outcome =
serde_json::to_value(data.outcome).map_err(|_| invalid("invalid outcome"))?;
let outcome = outcome.as_str().ok_or_else(|| invalid("invalid outcome"))?;
let end = at(data
.completed_at_ms
.ok_or_else(|| invalid("completed time required"))?)?;
let obs = serde_json::json!({"request":{"configRevision":1,"protocol":data.protocol,"endpoint":data.endpoint,"compact":false},"routing":{"scope":"all","groupRefs":[],"groupNamesSnapshot":[]}});
let (source, amount, currency) = cost_parts(&data.cost);
let mut q = QueryBuilder::new(
"insert into model_requests(id,client_api_key_ref,provider_kind,provider_account_ref,requested_model_id,upstream_model_id,operation,client_transport,service_tier,started_at,recovery_lease_expires_at,completed_at,outcome,request_observation_json,client_status_code,downstream_committed_at,input_tokens,output_tokens,cached_input_tokens,cache_write_input_tokens,reasoning_output_tokens,image_input_tokens,image_output_tokens,total_tokens,cost_source,cost_amount,cost_currency) values (",
);
q.push_bind(id)
.push(",")
.push_bind(&data.client_key_ref)
.push(",")
.push_bind(&data.provider)
.push(",")
.push_bind(&data.account_ref)
.push(",")
.push_bind(&data.requested_model)
.push(",")
.push_bind(&data.model)
.push(",")
.push_bind(&data.operation)
.push(",")
.push_bind(&data.client_transport)
.push(",")
.push_bind(&data.service_tier)
.push(",")
.push_bind(at(data.started_at_ms)?)
.push(",")
.push_bind(end)
.push(",")
.push_bind(end)
.push(",")
.push_bind(outcome)
.push(",")
.push_bind(obs)
.push(",")
.push_bind(if data.outcome == UsageOutcome::Succeeded {
Some(200i32)
} else {
None
})
.push(",")
.push_bind(if data.outcome == UsageOutcome::Succeeded {
Some(end)
} else {
None
});
for token in token_values(&data.tokens)? {
q.push(",").push_bind(token);
}
q.push(",")
.push_bind(source)
.push(",")
.push_bind(amount)
.push("::numeric,")
.push_bind(currency)
.push(") on conflict(id) do update set ");
for (i, column) in [
"client_api_key_ref",
"provider_kind",
"provider_account_ref",
"requested_model_id",
"upstream_model_id",
"operation",
"client_transport",
"service_tier",
"started_at",
"recovery_lease_expires_at",
"completed_at",
"outcome",
"request_observation_json",
"client_status_code",
"downstream_committed_at",
"input_tokens",
"output_tokens",
"cached_input_tokens",
"cache_write_input_tokens",
"reasoning_output_tokens",
"image_input_tokens",
"image_output_tokens",
"total_tokens",
"cost_source",
"cost_amount",
"cost_currency",
]
.iter()
.enumerate()
{
if i > 0 {
q.push(",");
}
q.push(format!("{column}=excluded.{column}"));
}
q.push(",billing_snapshot_json=null where model_requests.outcome<>'running'");
let changed = q
.build()
.execute(&mut **tx)
.await
.map_err(db_error)?
.rows_affected()
> 0;
if !changed {
return Err(invalid("cannot replace running usage record"));
}
Ok(true)
}
UsageRecordChange::Update { id, tokens, cost } => {
if running.is_none() {
return Err(StoreError::NotFound {
entity: "usage record",
id: id.clone(),
source: None,
});
}
let mut q = QueryBuilder::new("update model_requests set billing_snapshot_json=null");
if let Some(tokens) = tokens {
for (name, value) in [
"input_tokens",
"output_tokens",
"cached_input_tokens",
"cache_write_input_tokens",
"reasoning_output_tokens",
"image_input_tokens",
"image_output_tokens",
"total_tokens",
]
.into_iter()
.zip(token_values(tokens)?)
{
q.push(format!(",{name}=")).push_bind(value);
}
}
if let Some(cost) = cost {
let (source, amount, currency) = cost_parts(cost);
q.push(",cost_source=")
.push_bind(source)
.push(",cost_amount=")
.push_bind(amount)
.push("::numeric,cost_currency=")
.push_bind(currency);
}
q.push(" where id=").push_bind(id);
q.build().execute(&mut **tx).await.map_err(db_error)?;
Ok(true)
}
}
}
@@ -39,6 +39,7 @@ mod runtime_settings;
mod schema_integrity;
mod snapshot;
mod snapshots;
mod usage_data;
mod usage_paging;
static TEST_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations");
@@ -0,0 +1,599 @@
//! 通用用量维护的原生统计、重试、原子性、重建与资源预算验收
use super::{TestDatabase, drain_analytics, observability_query_budget, observability_repository};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::{
model::{
auth::{AdminAuditEvent, AuditActorKind},
usage_data::*,
},
ports::usage_data::*,
};
use gateway_store::postgres::{
ObservabilityQueryBudget, ObservabilityRange, PgRequestAnalyticsRepository, PgUsageDataStore,
UsageRecordFilter as NativeFilter,
};
use std::{collections::BTreeMap, time::Duration};
fn start() -> DateTime<Utc> {
DateTime::parse_from_rfc3339("2026-10-01T00:00:00Z")
.unwrap()
.with_timezone(&Utc)
}
fn filter() -> UsageRecordFilter {
UsageRecordFilter {
start_ms: start().timestamp_millis(),
end_ms: (start() + TimeDelta::days(1)).timestamp_millis(),
provider: None,
model: None,
account_ref: None,
client_key_ref: None,
}
}
fn audit() -> AdminAuditEvent {
AdminAuditEvent {
id: uuid::Uuid::new_v4().to_string(),
actor_kind: AuditActorKind::System,
actor_admin_user_id: None,
actor_ref: "test".into(),
request_id: None,
action: "usage.put".into(),
entity_kind: "usage".into(),
entity_ref: "test".into(),
config_revision: None,
changed_fields: vec!["usage".into()],
occurred_at: Utc::now(),
}
}
fn records(changes: Vec<UsageRecordChange>) -> StoredUsageWrite {
StoredUsageWrite {
changes: changes.into_iter().map(|c| (c, audit())).collect(),
}
}
fn aggregates(changes: Vec<UsageAggregateChange>) -> StoredAggregateWrite {
StoredAggregateWrite {
changes: changes.into_iter().map(|c| (c, audit())).collect(),
}
}
fn data(total: u64) -> UsageRecordData {
UsageRecordData {
client_key_ref: "key_usage".into(),
provider: Some("openai".into()),
account_ref: Some("account_usage".into()),
requested_model: Some("model_a".into()),
model: Some("model_a".into()),
protocol: "openai".into(),
operation: "responses".into(),
endpoint: "/v1/responses".into(),
client_transport: "http_sse".into(),
service_tier: None,
started_at_ms: start().timestamp_millis(),
completed_at_ms: Some((start() + TimeDelta::seconds(1)).timestamp_millis()),
outcome: UsageOutcome::Succeeded,
tokens: UsageTokens {
input_tokens: Some(UsageCount(total)),
total_tokens: Some(UsageCount(total)),
..Default::default()
},
cost: UsageCost::Unavailable,
}
}
fn put(id: &str, total: u64) -> UsageRecordChange {
UsageRecordChange::Put {
id: id.into(),
data: Box::new(data(total)),
}
}
fn contribution(key: &str, requests: i64, tokens: i64, cost: &str) -> UsageAggregateChange {
UsageAggregateChange::Put {
aggregate: UsageAggregate {
identity: UsageAggregateKey {
bucket_start_ms: start().timestamp_millis(),
key: key.into(),
dimensions: UsageDimensions {
provider: Some("openai".into()),
model: Some("model_a".into()),
account_ref: Some("account_usage".into()),
client_key_ref: Some("key_usage".into()),
..Default::default()
},
},
counts: BTreeMap::from([
(UsageMetric::RequestCount, UsageDelta(requests)),
(UsageMetric::CompletedRequestCount, UsageDelta(requests)),
(UsageMetric::TotalTokens, UsageDelta(tokens)),
]),
costs: BTreeMap::from([("USD".into(), cost.into())]),
},
}
}
#[tokio::test]
async fn detail_crud_is_retry_safe_atomic_and_updates_native_statistics() {
let Some(db) = TestDatabase::create("usage_detail_crud").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
store
.write_records(records(vec![put("a", 10), put("b", 20)]))
.await
.unwrap();
store
.write_records(records(vec![put("a", 10)]))
.await
.unwrap();
let first = store
.list_records(ListUsageRecords {
filter: filter(),
cursor: None,
limit: Some(1),
})
.await
.unwrap();
let second = store
.list_records(ListUsageRecords {
filter: filter(),
cursor: first.next_cursor,
limit: Some(1),
})
.await
.unwrap();
assert_eq!(first.items[0].id, "a");
assert_eq!(second.items[0].id, "b");
assert!(second.next_cursor.is_none());
let mut invalid = data(50);
invalid.cost = UsageCost::Known {
source: UsageCostSource::Calculated,
amount: "invalid".into(),
currency: "USD".into(),
};
let err = store
.write_records(records(vec![
put("a", 99),
UsageRecordChange::Put {
id: "bad".into(),
data: Box::new(invalid),
},
]))
.await
.unwrap_err();
assert_eq!(err.kind, UsageErrorKind::InvalidInput);
assert_eq!(
store
.get_record(GetUsageRecord { id: "a".into() })
.await
.unwrap()
.data
.tokens
.total_tokens,
Some(UsageCount(10))
);
drain_analytics(&db.pool).await;
let repo = observability_repository(&db.pool);
let range = ObservabilityRange::new(start(), start() + TimeDelta::days(1)).unwrap();
let summary = repo
.usage_summary(range, NativeFilter::default())
.await
.unwrap();
assert_eq!(
(
summary.requests.request_count,
summary.requests.total_tokens
),
(2, 30)
);
store
.write_records(records(vec![
UsageRecordChange::Update {
id: "a".into(),
tokens: Some(data(100).tokens),
cost: Some(UsageCost::Known {
source: UsageCostSource::Calculated,
amount: "0".into(),
currency: "USD".into(),
}),
},
UsageRecordChange::Delete { id: "b".into() },
]))
.await
.unwrap();
drain_analytics(&db.pool).await;
let row = store
.get_record(GetUsageRecord { id: "a".into() })
.await
.unwrap();
assert!(matches!(row.data.cost, UsageCost::Known { .. }));
assert!(row.billing_snapshot.is_none());
let summary = repo
.usage_summary(range, NativeFilter::default())
.await
.unwrap();
assert_eq!(
(
summary.requests.request_count,
summary.requests.total_tokens
),
(1, 100)
);
assert!(
store
.get_record(GetUsageRecord { id: "b".into() })
.await
.is_err()
);
assert!(
store
.write_records(records(vec![UsageRecordChange::Delete { id: "b".into() }]))
.await
.unwrap()
.changed_ids
.is_empty()
);
db.close().await;
}
#[tokio::test]
async fn aggregate_import_and_signed_correction_survive_native_rebuild_without_schema_changes() {
let Some(db) = TestDatabase::create("usage_aggregate_rebuild").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
store
.write_records(records(vec![put("native", 120)]))
.await
.unwrap();
drain_analytics(&db.pool).await;
let imported = contribution("daily", 100, 1100, "2.75");
store
.write_aggregates(aggregates(vec![imported.clone()]))
.await
.unwrap();
store
.write_aggregates(aggregates(vec![
imported,
contribution("correction", -20, -100, "-1"),
]))
.await
.unwrap();
let repo = observability_repository(&db.pool);
let range = ObservabilityRange::new(start(), start() + TimeDelta::days(1)).unwrap();
for _ in 0..2 {
let summary = repo
.usage_summary(range, NativeFilter::default())
.await
.unwrap();
assert_eq!(
(
summary.requests.request_count,
summary.requests.total_tokens
),
(81, 1120)
);
let filtered = repo
.usage_summary(
range,
NativeFilter {
provider_account_ref: Some("account_usage".into()),
provider_kind: Some("openai".into()),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(filtered.requests.request_count, 81);
use gateway_admin::{
model::{accounts::AccountUsageWindowQuery, observability::TimeRange},
ports::store::AccountStore,
};
let windows = [AccountUsageWindowQuery {
account_id: "account_usage".into(),
key: "window".into(),
range: TimeRange {
start: range.start,
end: range.end,
},
}];
let usage = super::admin_account_store(&db.pool)
.load_account_usage_by_windows(&windows)
.await
.unwrap();
assert_eq!(usage[0].usage.request_count, 81);
assert_eq!(usage[0].usage.total_tokens, Some(1120));
assert_eq!(usage[0].usage.costs[0].amount.as_str(), "1.75");
let costs = repo
.usage_diagnostics(
range,
NativeFilter::default(),
gateway_store::postgres::DiagnosticDimension::Account,
10,
)
.await
.unwrap();
assert_eq!(costs.total_request_count, 81);
assert_eq!(costs.items[0].costs[0].amount.as_str(), "1.75");
sqlx::query(
"update request_analytics_build_state set schema_version=0,histogram_version=0",
)
.execute(&db.pool)
.await
.unwrap();
drain_analytics(&db.pool).await;
}
let page = store
.list_aggregates(ListUsageAggregates {
filter: filter(),
cursor: None,
limit: Some(1),
})
.await
.unwrap();
let next = store
.list_aggregates(ListUsageAggregates {
filter: filter(),
cursor: page.next_cursor,
limit: Some(1),
})
.await
.unwrap();
assert_eq!(page.items.len() + next.items.len(), 2);
let count: i64 = sqlx::query_scalar("select count(*) from model_requests")
.fetch_one(&db.pool)
.await
.unwrap();
assert_eq!(count, 1, "aggregates must not fabricate request details");
// 显式桶重建只重建请求投影,不能删除手工维护的统计贡献
sqlx::query("insert into request_analytics_pending_buckets(bucket_start,oldest_pending_at) values($1,now()) on conflict(bucket_start) do update set change_version=nextval('request_analytics_change_version')").bind(start()).execute(&db.pool).await.unwrap();
let worker = PgRequestAnalyticsRepository::new(db.pool.clone(), observability_query_budget());
for _ in 0..30 {
worker.maintain(500).await.unwrap();
}
assert_eq!(
repo.usage_summary(range, NativeFilter::default())
.await
.unwrap()
.requests
.request_count,
81
);
store
.write_aggregates(aggregates(
page.items
.into_iter()
.chain(next.items)
.map(|a| UsageAggregateChange::Delete {
identity: a.identity,
})
.collect(),
))
.await
.unwrap();
assert_eq!(
repo.usage_summary(range, NativeFilter::default())
.await
.unwrap()
.requests
.request_count,
1
);
db.close().await;
}
#[tokio::test]
async fn empty_native_history_still_shows_aggregates_and_bad_values_return_errors() {
let Some(db) = TestDatabase::create("usage_aggregate_bounds").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
drain_analytics(&db.pool).await;
store
.write_aggregates(aggregates(vec![contribution("only", 7, 90, "1")]))
.await
.unwrap();
let repo = observability_repository(&db.pool);
let range = ObservabilityRange::new(start(), start() + TimeDelta::days(1)).unwrap();
assert_eq!(
repo.usage_summary(range, NativeFilter::default())
.await
.unwrap()
.requests
.request_count,
7
);
store
.write_aggregates(aggregates(vec![contribution("only", i64::MIN, 90, "1")]))
.await
.unwrap();
assert!(
repo.usage_summary(range, NativeFilter::default())
.await
.is_err()
);
store
.write_aggregates(aggregates(vec![contribution("only", 8, 90, "1")]))
.await
.unwrap();
assert_eq!(
repo.usage_summary(range, NativeFilter::default())
.await
.unwrap()
.requests
.request_count,
8
);
db.close().await;
}
#[tokio::test]
async fn timeout_and_cancellation_release_connections_and_rollback_the_whole_batch() {
let Some(db) = TestDatabase::create("usage_budget").await else {
return;
};
let budget = ObservabilityQueryBudget::try_new(
1,
Duration::from_millis(200),
Duration::from_millis(300),
)
.unwrap();
let store = PgUsageDataStore::new(db.pool.clone(), budget);
store
.write_records(records(vec![put("locked", 10)]))
.await
.unwrap();
let mut lock = db.pool.begin().await.unwrap();
sqlx::query("select id from model_requests where id='locked' for update")
.fetch_one(&mut *lock)
.await
.unwrap();
assert_eq!(
store
.write_records(records(vec![put("first", 9), put("locked", 20)]))
.await
.unwrap_err()
.kind,
UsageErrorKind::Unavailable
);
lock.rollback().await.unwrap();
assert_eq!(
store
.get_record(GetUsageRecord { id: "first".into() })
.await
.unwrap_err()
.kind,
UsageErrorKind::NotFound
);
assert_eq!(
store
.get_record(GetUsageRecord {
id: "locked".into()
})
.await
.unwrap()
.data
.tokens
.total_tokens,
Some(UsageCount(10))
);
let mut lock = db.pool.begin().await.unwrap();
sqlx::query("select id from model_requests where id='locked' for update")
.fetch_one(&mut *lock)
.await
.unwrap();
let cloned = store.clone();
let task = tokio::spawn(async move {
cloned
.write_records(records(vec![put("cancelled", 9), put("locked", 30)]))
.await
});
tokio::time::sleep(Duration::from_millis(30)).await;
task.abort();
let _ = task.await;
lock.rollback().await.unwrap();
tokio::time::sleep(Duration::from_millis(30)).await;
assert_eq!(
store
.get_record(GetUsageRecord {
id: "cancelled".into()
})
.await
.unwrap_err()
.kind,
UsageErrorKind::NotFound
);
store
.write_records(records(vec![put("after", 1)]))
.await
.unwrap();
db.close().await;
}
#[tokio::test]
async fn unassigned_correction_changes_totals_without_inventing_dimension_members() {
let Some(db) = TestDatabase::create("usage_global_correction").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
store
.write_records(records(vec![put("a", 100), put("b", 100)]))
.await
.unwrap();
drain_analytics(&db.pool).await;
let mut correction = contribution("global", -1, -50, "0");
let UsageAggregateChange::Put { aggregate } = &mut correction else {
unreachable!()
};
aggregate.identity.dimensions = UsageDimensions::default();
store
.write_aggregates(aggregates(vec![correction]))
.await
.unwrap();
let repo = observability_repository(&db.pool);
let range = ObservabilityRange::new(start(), start() + TimeDelta::days(1)).unwrap();
let summary = repo
.usage_summary(range, NativeFilter::default())
.await
.unwrap();
assert_eq!(
(
summary.requests.request_count,
summary.requests.total_tokens
),
(1, 150)
);
assert_eq!(summary.providers.len(), 1);
assert_eq!(summary.providers[0].request_count, 2);
db.close().await;
}
#[tokio::test]
async fn aggregate_corrections_participate_in_bounded_account_and_key_ranking() {
let Some(db) = TestDatabase::create("usage_aggregate_ranking").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
let mut other = data(10);
other.account_ref = Some("account_other".into());
other.client_key_ref = "key_other".into();
store
.write_records(records(vec![
put("a", 10),
put("b", 10),
UsageRecordChange::Put {
id: "c".into(),
data: Box::new(other),
},
]))
.await
.unwrap();
drain_analytics(&db.pool).await;
let mut imported = contribution("other", 10, 100, "0");
let UsageAggregateChange::Put { aggregate } = &mut imported else {
unreachable!()
};
aggregate.identity.dimensions.account_ref = Some("account_other".into());
aggregate.identity.dimensions.client_key_ref = Some("key_other".into());
store
.write_aggregates(aggregates(vec![
imported,
contribution("correction", -1, -10, "0"),
]))
.await
.unwrap();
let repo = observability_repository(&db.pool);
let range = ObservabilityRange::new(start(), start() + TimeDelta::days(1)).unwrap();
use gateway_store::postgres::DiagnosticDimension;
for (dimension, key) in [
(DiagnosticDimension::Account, "account_other"),
(DiagnosticDimension::ApiKey, "key_other"),
] {
let result = repo
.usage_diagnostics(range, NativeFilter::default(), dimension, 1)
.await
.unwrap();
assert_eq!(result.total_request_count, 12);
assert_eq!(result.items.len(), 1);
assert_eq!(result.items[0].key, key);
assert_eq!(result.items[0].request_count, 11);
}
db.close().await;
}
@@ -433,6 +433,33 @@ impl ProviderAdmin for OpenAiAdminProvider {
}))
}
fn evaluate_usage_pricing(
&self,
input: &gateway_admin::ports::provider::ProviderUsagePricingInput,
pricing: Option<&gateway_core::metering::ModelPriceOverride>,
) -> Result<Option<gateway_core::metering::CalculatedCostBreakdown>, ProviderAdminError> {
let (Some(input_tokens), Some(output_tokens), Some(cached), Some(written)) = (
input.usage.input_tokens,
input.usage.output_tokens,
input.usage.cached_input_tokens,
input.usage.cache_write_input_tokens,
) else {
return Ok(None);
};
if input.usage.image_input_tokens != Some(0) || input.usage.image_output_tokens != Some(0) {
return Ok(None);
}
let usage = OpenAiBillingUsage::new(input_tokens, output_tokens, cached, written);
Ok(
crate::transport::usage::openai_billing_breakdown_with_override(
&input.upstream_model_id,
usage,
input.service_tier.as_deref(),
pricing,
),
)
}
async fn prepare_import(
&self,
command: PrepareCredentialImport,
@@ -2776,3 +2776,40 @@ impl gateway_core::diagnostics::OperationalDiagnostics for RecordingDiagnostics
Ok(())
}
}
#[tokio::test]
async fn usage_quote_reuses_native_override_and_reports_incomplete_input() {
use gateway_admin::ports::provider::ProviderUsagePricingInput;
use gateway_core::metering::{ModelPriceOverride, Usage};
let config = valid_config();
let bundle = provider_openai::initialize(config.config.clone(), provider_ports())
.await
.unwrap();
let admin = bundle.admin_provider();
let mut input = ProviderUsagePricingInput {
upstream_model_id: "gpt-6-sol".into(),
service_tier: None,
usage: Usage {
input_tokens: Some(1_000_000),
output_tokens: Some(100_000),
cached_input_tokens: Some(200_000),
cache_write_input_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
..Default::default()
},
};
let custom:ModelPriceOverride=serde_json::from_value(json!({"multiplierBps":10000,"bands":{"standard":{"input":"2","output":"10","cacheRead":"0.5","cacheWrite":"2"},"long_standard":{"input":"2","output":"10","cacheRead":"0.5","cacheWrite":"2"}}})).unwrap();
let price = admin
.evaluate_usage_pricing(&input, Some(&custom))
.unwrap()
.unwrap();
assert_eq!(price.total_amount().amount().canonical(), "2.7");
input.usage.cached_input_tokens = None;
assert!(
admin
.evaluate_usage_pricing(&input, Some(&custom))
.unwrap()
.is_none()
);
}
+29
View File
@@ -489,6 +489,35 @@ impl ProviderAdmin for XaiAdminProvider {
}))
}
fn evaluate_usage_pricing(
&self,
input: &gateway_admin::ports::provider::ProviderUsagePricingInput,
pricing: Option<&gateway_core::metering::ModelPriceOverride>,
) -> Result<Option<gateway_core::metering::CalculatedCostBreakdown>, ProviderAdminError> {
let (Some(input_tokens), Some(output_tokens), Some(cached), Some(written)) = (
input.usage.input_tokens,
input.usage.output_tokens,
input.usage.cached_input_tokens,
input.usage.cache_write_input_tokens,
) else {
return Ok(None);
};
if input.usage.image_input_tokens != Some(0) || input.usage.image_output_tokens != Some(0) {
return Ok(None);
}
Ok(
crate::transport::canonical::grok_billing_breakdown_with_override(
&input.upstream_model_id,
input_tokens,
output_tokens,
cached,
written,
input.service_tier.as_deref(),
pricing,
),
)
}
async fn prepare_import(
&self,
command: PrepareCredentialImport,
@@ -1385,3 +1385,26 @@ impl gateway_core::diagnostics::OperationalDiagnostics for RecordingDiagnostics
Ok(())
}
}
#[tokio::test]
async fn usage_quote_reuses_native_priority_prices() {
use gateway_admin::ports::provider::ProviderUsagePricingInput;
use gateway_core::metering::Usage;
let bundle = provider_xai::initialize(provider_ports()).await.unwrap();
let admin = bundle.admin_provider();
let input = ProviderUsagePricingInput {
upstream_model_id: "grok-4.5".into(),
service_tier: Some("priority".into()),
usage: Usage {
input_tokens: Some(100),
output_tokens: Some(10),
cached_input_tokens: Some(25),
cache_write_input_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
..Default::default()
},
};
let price = admin.evaluate_usage_pricing(&input, None).unwrap().unwrap();
assert_eq!(price.total_amount().amount().canonical(), "0.000435");
}
+2
View File
@@ -2106,6 +2106,8 @@ Host 关闭或任务取消会记录失败终态;状态查询会收敛无执行
配置完整则启用,缺少必填项则保留为待配置。操作流程见 [插件使用](plugins.md),开发合同见
[SDK](../backend/crates/gateway-plugin/sdk/README.md)。不提供官方商店或第三方市场订阅接口
插件通过类型化公开服务维护请求明细、聚合统计并试算费用,接口与数值合同见 [SDK 用量维护与计价](../backend/crates/gateway-plugin/sdk/docs/usage-data.md)。聚合参与原生总量与趋势,不增加请求明细列表条目
### 安装包与来源
| 方法 | 路径 | 行为 |
+8
View File
@@ -773,6 +773,14 @@ PostgreSQL 不可用时拒绝所有新的计费请求,Redis 继续管理并发
账本独立于可丢弃的请求观测日志,日志清理不重置金额;费用事件保留至删除 Key,已有日志不会回填为账本费用。
字段与错误合同见 [Client Key API](api.md#7-client-key)
### 用量维护服务
Admin 定义明细、聚合维护与纯计价合同,Provider 复用原生 Token 计价公式,Runtime 传递可信调用身份。插件负责外部格式、来源映射、去重、调整历史、目标选择和跨调用流程
明细直接维护 `model_requests` 的终态记录,沿用现有统计失效触发器;计价快照在维护后清空,Key 结算账本保持独立。聚合贡献复用 `request_analytics_1h` 的 `collection`、维度、计数和费用结构,`usage_contribution` 集合保存调用方输入,其他集合仍为请求投影。该集合独立于请求投影,桶重建和版本重建保留该集合,查询在同一快照中按指定维度和归属时间合并贡献;聚合的删除由调用方管理
Store 在单次事务内完成有限批量读写和普通审计,复用观测连接与执行预算,不增加数据库表或字段。原生统计结果仍遵守现有数值类型,越界返回错误,不自动修正插件数据。调用合同与时间归属见 [SDK 用量维护与计价](../backend/crates/gateway-plugin/sdk/docs/usage-data.md)
### 模型价格与费用快照
Provider 是内置价目、服务档位、模态和工具费规则的唯一 owner;Core metering 定义中立价格覆盖与
+3 -1
View File
@@ -139,10 +139,12 @@ flowchart LR
| 停用 | 不删除;新请求不再使用 | 配置、密钥、私有数据和版本 |
| 删除配置 | 该配置及其密钥、私有数据、版本配置快照;无法撤销 | 已安装版本 |
| 删除版本 | 所选制品、对应配置快照及其已无引用的下载认证;最后一个版本还会清理来源规则;被配置或在途调用引用时不能删除 | 其他版本与配置、共享下载认证 |
| 卸载插件 | 确认范围内的配置、数据、版本、来源规则及已无引用的下载认证 | 其他插件仍使用的下载认证、审计记录 |
| 卸载插件 | 确认范围内的配置、插件私有数据、版本、来源规则及已无引用的下载认证 | 其他插件仍使用的下载认证、审计记录 |
卸载不是跨多个请求的事务;失败时会显示已完成结果,可重试剩余步骤,不会扩大到确认后新增的配置。只想暂时停用时不要卸载。
完整卸载后重新安装可重新选择来源。独立创建但未被删除版本使用的下载认证不会被顺带删除。
通过用量服务写入的明细与聚合统计不会随插件停用、删除配置或卸载而删除。明细遵循宿主请求历史保留策略,聚合由写入方管理清理;删除插件前按需使用其数据清理功能
需要保留数据时,删除前按[部署、备份与恢复](../deploy/README.md)保存备份
开发插件见 [SDK](../backend/crates/gateway-plugin/sdk/README.md) 与 [插件 CLI](../backend/apps/plugin-cli/README.md);