feat(usage): add snapshot scans and complete grouped statistics (#599)

Expose bounded usage readers through Admin, Store, Runtime, and SDK. Preserve snapshot consistency, coverage metadata, cancellation, and parent-call cleanup.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-11-04 15:22:44 +08:00
committed by GitHub
co-authored by Zyy Codex
parent 0663df14df
commit 78e943f406
44 changed files with 2471 additions and 98 deletions
@@ -146,6 +146,10 @@ async fn execute(
providers.admin,
);
or_shutdown!(plugin_runtime, services.register_usage_data(&usage_data));
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_usage_queries(&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);
@@ -160,6 +160,10 @@ pub async fn run() -> Result<(), BootstrapError> {
plugin_runtime,
plugin_runtime.bind_services(&admin.services().public_services())
);
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_usage_queries(&admin.services().usage_data())
);
let mut plan = store.take_worker_contributions();
plan.push(analytics_worker);
+2 -2
View File
@@ -288,8 +288,8 @@ impl AdminServices {
}
#[must_use]
pub fn usage_data(&self) -> &dyn UsageDataService {
self.usage_data.as_ref()
pub fn usage_data(&self) -> Arc<dyn UsageDataService> {
self.usage_data.clone()
}
#[must_use]
@@ -10,6 +10,8 @@ pub enum UsageErrorKind {
Conflict,
CapacityExceeded,
Unavailable,
Cancelled,
Timeout,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -20,3 +20,6 @@ pub use actor::*;
mod aggregate;
pub use aggregate::*;
mod read;
pub use read::*;
@@ -0,0 +1,162 @@
//! 历史用量的窄事实、分组统计与调用内分批读取合同
use super::{UsageCount, UsageMetric, UsageRecordData, UsageRecordFilter};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
pub const USAGE_READ_MAX_BATCH_ROWS: u16 = 1024;
pub const USAGE_READ_MAX_BATCH_BYTES: usize = 2 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadFilter {
pub records: UsageRecordFilter,
pub routing_group_ref: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScanUsageRecords {
pub filter: UsageReadFilter,
pub batch_size: Option<u16>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageGrouping {
Model,
Key,
KeyModel,
RoutingGroup,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct QueryUsageStatistics {
pub filter: UsageReadFilter,
pub group_by: UsageGrouping,
pub batch_size: Option<u16>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum UsageReadQuery {
Records(ScanUsageRecords),
Statistics(QueryUsageStatistics),
}
impl UsageReadQuery {
pub fn filter(&self) -> &UsageReadFilter {
match self {
Self::Records(query) => &query.filter,
Self::Statistics(query) => &query.filter,
}
}
pub fn batch_size(&self) -> u16 {
match self {
Self::Records(query) => query.batch_size,
Self::Statistics(query) => query.batch_size,
}
.unwrap_or(USAGE_READ_MAX_BATCH_ROWS)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageRecordFact {
pub id: String,
pub data: UsageRecordData,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
pub subagent_kind: Option<String>,
pub request_kind: Option<String>,
pub routing_scope: Option<String>,
pub routing_group_refs: Vec<String>,
pub recovered: bool,
/// 由宿主现行规则派生,原始审计中的恢复请求仍可读取
pub request_counted: bool,
/// 包含完整交付、预热排除及恢复关联规则
pub usage_counted: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsGroup {
pub client_key_ref: Option<String>,
pub model: Option<String>,
pub routing_group_ref: Option<String>,
pub unassigned_contribution: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsCost {
pub currency: String,
pub amount: String,
pub native_amount: Option<String>,
/// 外部修正允许负值,原生费用来源覆盖单独返回
pub contribution_amount: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsCostCoverage {
pub provider_reported: UsageCount,
pub calculated: UsageCount,
pub unavailable: UsageCount,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsMetrics {
/// 缺失的 Token 指标表示未知,计数指标始终返回
pub values: BTreeMap<UsageMetric, UsageCount>,
pub native_request_count: UsageCount,
pub contribution_count: UsageCount,
pub native_token_samples: BTreeMap<UsageMetric, UsageCount>,
/// 数量表示提供该指标的贡献记录数,不能用作请求覆盖率的分母
pub contribution_token_samples: BTreeMap<UsageMetric, UsageCount>,
pub cost_coverage: UsageStatisticsCostCoverage,
pub costs: Vec<UsageStatisticsCost>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsRow {
/// null 表示全局总计,分组重叠不会重复计入该行
pub group: Option<UsageStatisticsGroup>,
pub metrics: UsageStatisticsMetrics,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
pub enum UsageReadItem {
Record(Box<UsageRecordFact>),
Statistics(Box<UsageStatisticsRow>),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadAnalytics {
pub coverage_start_ms: Option<i64>,
pub coverage_end_ms: Option<i64>,
pub lag_seconds: u64,
pub state: String,
pub generation: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadObservation {
pub generated_at_ms: i64,
pub analytics: Option<UsageReadAnalytics>,
pub total_groups: Option<UsageCount>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadBatch {
pub observation: UsageReadObservation,
pub items: Vec<UsageReadItem>,
pub complete: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadOpened {
pub reader: Option<String>,
pub batch: UsageReadBatch,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadHandle {
pub reader: String,
}
@@ -2,6 +2,14 @@
use crate::model::{auth::AdminAuditEvent, usage_data::*};
use async_trait::async_trait;
use gateway_core::lifecycle::CancellationToken;
#[async_trait]
pub trait UsageReadSession: Send {
async fn next_batch(&mut self) -> UsageDataResult<UsageReadBatch>;
fn is_finished(&self) -> bool;
async fn close(&mut self);
}
pub struct StoredUsageWrite {
pub changes: Vec<(UsageRecordChange, AdminAuditEvent)>,
@@ -13,6 +21,11 @@ pub struct StoredAggregateWrite {
#[async_trait]
pub trait UsageDataStore: Send + Sync {
async fn open_read(
&self,
query: UsageReadQuery,
cancellation: CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>>;
async fn list_aggregates(
&self,
request: ListUsageAggregates,
@@ -15,6 +15,11 @@ use std::sync::Arc;
#[async_trait]
pub trait UsageDataService: Send + Sync {
async fn open_read(
&self,
query: UsageReadQuery,
cancellation: gateway_core::lifecycle::CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>>;
async fn list_aggregates(
&self,
request: ListUsageAggregates,
@@ -58,6 +63,21 @@ impl DefaultUsageDataService {
#[async_trait]
impl UsageDataService for DefaultUsageDataService {
async fn open_read(
&self,
query: UsageReadQuery,
cancellation: gateway_core::lifecycle::CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>> {
let filter = query.filter();
validation::query(&filter.records, None, None)?;
if let Some(group) = &filter.routing_group_ref {
validation::text(group, 128)?;
}
if !(1..=USAGE_READ_MAX_BATCH_ROWS).contains(&query.batch_size()) {
return Err(UsageDataError::invalid("读取批次条数应为 1 到 1024"));
}
self.store.open_read(query, cancellation).await
}
async fn list_aggregates(
&self,
request: ListUsageAggregates,
@@ -3,6 +3,14 @@ use super::TestPluginPorts;
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
#[async_trait::async_trait]
impl UsageDataStore for TestPluginPorts {
async fn open_read(
&self,
_: UsageReadQuery,
_: gateway_core::lifecycle::CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>> {
Err(UsageDataError::unavailable("test usage reader unavailable"))
}
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
unreachable!("unexpected usage call")
}
@@ -64,6 +64,7 @@ fn sdk_wire_types_match_host_contracts() {
"value.rs",
"record.rs",
"query.rs",
"read.rs",
"aggregate.rs",
"pricing.rs",
"error.rs",
@@ -3,6 +3,14 @@ use super::TestPluginPorts;
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
#[async_trait::async_trait]
impl UsageDataStore for TestPluginPorts {
async fn open_read(
&self,
_: UsageReadQuery,
_: gateway_core::lifecycle::CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>> {
Err(UsageDataError::unavailable("test usage reader unavailable"))
}
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
unreachable!("unexpected usage call")
}
@@ -110,6 +110,10 @@ impl CallResources {
stream.close();
}
state.model_streams.clear();
for reader in state.usage_readers.values() {
reader.cancel();
}
state.usage_readers.clear();
}
pub(super) fn start_upstream(
@@ -179,6 +183,7 @@ impl CallResources {
#[derive(Default)]
pub(super) struct CallState {
pub(super) usage_readers: BTreeMap<String, Arc<super::usage::UsageReader>>,
pub(super) closed: bool,
pub(super) middleware: Option<Arc<dyn MiddlewareCallback>>,
pub(super) streams: BTreeMap<String, Arc<HttpStream>>,
@@ -79,6 +79,7 @@ impl CallbackHandler for PluginCallbacks {
let request_limits = self.request_limits.clone();
let affinity = self.affinity.clone();
let services = self.services.clone();
let usage = self.usage.clone();
let service_actor = self.service_actor.clone();
Box::pin(async move {
let call = call.ok_or_else(denied)?;
@@ -150,6 +151,15 @@ impl CallbackHandler for PluginCallbacks {
.call(&context, &call, service_actor, params, &payload)
.await;
}
if matches!(
method.as_str(),
gateway_plugin_sdk::call::services::usage::RECORDS_SCAN
| gateway_plugin_sdk::call::services::usage::STATISTICS_QUERY
| gateway_plugin_sdk::call::services::usage::READ_NEXT
| gateway_plugin_sdk::call::services::usage::READ_CLOSE
) {
return usage.call(&call, &method, params, &payload).await;
}
if is_private_state {
return private_state.call(&method, params, &payload).await;
}
@@ -26,6 +26,7 @@ mod resources;
mod scope;
pub(crate) mod services;
pub(crate) mod upstream;
pub(crate) mod usage;
pub(crate) mod websocket_middleware;
pub(crate) use accounts::PluginAccountPortSlot;
@@ -19,6 +19,7 @@ use super::{
};
pub(crate) struct PluginCallbackPorts {
pub(crate) usage: Arc<super::usage::UsagePorts>,
pub(crate) services: Arc<services::ServicePorts>,
pub(crate) http: Arc<HttpClient>,
pub(crate) accounts: Arc<PluginAccountPortSlot>,
@@ -30,6 +31,7 @@ pub(crate) struct PluginCallbackPorts {
}
pub(crate) struct PluginCallbacks {
pub(super) usage: Arc<super::usage::UsagePorts>,
pub(crate) registration_body_limit: Option<usize>,
pub(super) services: Arc<services::ServicePorts>,
pub(super) service_actor: gateway_admin::model::usage_data::UsageActor,
@@ -63,6 +65,7 @@ impl PluginCallbacks {
) -> Result<Self, AdminError> {
let clients = Arc::new(client::PluginClients::new(ports.clients));
Ok(Self {
usage: ports.usage,
registration_body_limit: None,
services: ports.services,
service_actor: gateway_admin::model::usage_data::UsageActor::Plugin {
@@ -0,0 +1,152 @@
//! 将用量读取绑定到父调用,Admin 拥有查询语义,Runtime 管理句柄和回收
use super::{CallResources, denied, invalid};
use crate::RpcReply;
use gateway_admin::{
UsageDataService,
model::{AdminError, usage_data::*},
ports::usage_data::UsageReadSession,
};
use gateway_core::lifecycle::CancellationToken;
use gateway_plugin_sdk::{ErrorCode, PluginFault, call::services::usage as wire};
use std::sync::{Arc, OnceLock, Weak};
use tokio::sync::Mutex;
#[derive(Default)]
pub(crate) struct UsagePorts(OnceLock<Weak<dyn UsageDataService>>);
pub(super) struct UsageReader {
session: Mutex<Box<dyn UsageReadSession>>,
cancellation: CancellationToken,
}
impl UsageReader {
pub(super) fn cancel(&self) {
self.cancellation.cancel();
}
}
impl UsagePorts {
pub(crate) fn bind(&self, service: &Arc<dyn UsageDataService>) -> Result<(), AdminError> {
self.0
.set(Arc::downgrade(service))
.map_err(|_| AdminError::conflict("用量查询端口已经绑定"))
}
pub(super) async fn call(
&self,
call: &CallResources,
method: &str,
params: serde_json::Value,
payload: &[u8],
) -> Result<RpcReply, PluginFault> {
if params != serde_json::json!({}) || payload.len() > USAGE_READ_MAX_BATCH_BYTES {
return Err(invalid());
}
match method {
wire::RECORDS_SCAN | wire::STATISTICS_QUERY => {
let query = if method == wire::RECORDS_SCAN {
UsageReadQuery::Records(serde_json::from_slice(payload).map_err(|_| invalid())?)
} else {
UsageReadQuery::Statistics(
serde_json::from_slice(payload).map_err(|_| invalid())?,
)
};
let service = self
.0
.get()
.and_then(Weak::upgrade)
.ok_or_else(|| PluginFault::new(ErrorCode::Fault, "用量查询端口不可用"))?;
let cancellation = call.cancellation.child_token();
let mut session = service
.open_read(query, cancellation.clone())
.await
.map_err(fault)?;
let batch = session.next_batch().await.map_err(fault)?;
let reader = (!batch.complete).then(|| uuid::Uuid::new_v4().to_string());
let reply = super::admin::encode(&UsageReadOpened {
reader: reader.clone(),
batch,
})?;
if let Some(reader) = reader {
let mut state = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.closed {
return Err(denied());
}
state.usage_readers.retain(|_, reader| {
!reader
.session
.try_lock()
.is_ok_and(|session| session.is_finished())
});
state.usage_readers.insert(
reader,
Arc::new(UsageReader {
session: Mutex::new(session),
cancellation,
}),
);
}
Ok(reply)
}
wire::READ_NEXT | wire::READ_CLOSE => {
let request: UsageReadHandle =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let reader = {
let mut state = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if method == wire::READ_CLOSE {
state.usage_readers.remove(&request.reader)
} else {
state.usage_readers.get(&request.reader).cloned()
}
}
.ok_or_else(|| {
fault(UsageDataError::new(
UsageErrorKind::NotFound,
"用量读取句柄无效或已结束",
))
})?;
if method == wire::READ_CLOSE {
reader.cancel();
reader.session.lock().await.close().await;
return super::admin::encode(&());
}
let mut session = reader.session.try_lock().map_err(|_| {
PluginFault::new(ErrorCode::Conflict, "同一用量读取只允许一个消费者")
})?;
let result = session.next_batch().await;
if result.as_ref().map_or(true, |batch| batch.complete) {
call.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.usage_readers
.remove(&request.reader);
session.close().await;
}
super::admin::encode(&result.map_err(fault)?)
}
_ => Err(invalid()),
}
}
}
fn fault(error: UsageDataError) -> PluginFault {
let code = match error.kind {
UsageErrorKind::InvalidInput => ErrorCode::InvalidInput,
UsageErrorKind::NotFound => ErrorCode::Conflict,
UsageErrorKind::Conflict => ErrorCode::Conflict,
UsageErrorKind::CapacityExceeded => ErrorCode::Capacity,
UsageErrorKind::Unavailable => ErrorCode::Fault,
UsageErrorKind::Cancelled => ErrorCode::Cancelled,
UsageErrorKind::Timeout => ErrorCode::Timeout,
};
let mut fault = PluginFault::new(code, error.message.clone());
fault.details = serde_json::to_value(error).ok();
fault
}
@@ -192,6 +192,7 @@ impl PluginRuntime {
private_state.clone(),
PluginCallbackPorts {
services: self.service_ports.clone(),
usage: self.usage_ports.clone(),
http: self.http.clone(),
accounts: self.account_ports.clone(),
keys: self.client_key_ports.clone(),
@@ -51,6 +51,7 @@ pub struct PluginRuntimeConfig {
/// 只索引有主人的候选;当前代次仅由 Core 的发布视图持有
pub struct PluginRuntime {
pub(super) usage_ports: Arc<crate::callback::usage::UsagePorts>,
pub(super) request_limits: Arc<crate::callback::request_limits::RequestLimits>,
pub(super) service_ports: Arc<crate::callback::services::ServicePorts>,
pub(super) store: Arc<dyn PluginStore>,
@@ -91,6 +92,7 @@ impl PluginRuntime {
) -> Self {
let restart_circuit = config.restart_circuit;
Self {
usage_ports: Arc::default(),
store,
state,
log_slots: Arc::new(Semaphore::new(
@@ -156,6 +158,13 @@ impl PluginRuntime {
self.service_ports.bind(registry)
}
pub fn bind_usage_queries(
&self,
service: &Arc<dyn gateway_admin::UsageDataService>,
) -> Result<(), AdminError> {
self.usage_ports.bind(service)
}
pub fn bind_account_ports(
&self,
access: &Arc<dyn PluginAccountAccess>,
@@ -472,6 +472,8 @@ async fn usage_sdk_calls_preserve_exact_counts_and_write_actor() {
let admin = environment.bind_admin_accounts(&runtime, &core).await;
let services = admin.services();
runtime.bind_services(&services.public_services()).unwrap();
runtime.bind_usage_queries(&services.usage_data()).unwrap();
environment.prepare_empty_analytics().await;
assert_eq!(
services
.public_services()
@@ -502,3 +504,117 @@ async fn usage_sdk_calls_preserve_exact_counts_and_write_actor() {
runtime.shutdown().await;
}
#[tokio::test]
async fn usage_reader_is_released_when_parent_call_is_cancelled() {
usage_reader_lifecycle("cancel").await;
}
#[tokio::test]
async fn disabled_instance_keeps_inflight_usage_read_until_parent_completion() {
usage_reader_lifecycle("disable").await;
}
#[tokio::test]
async fn quiescing_instance_releases_inflight_usage_read() {
usage_reader_lifecycle("quiesce").await;
}
async fn usage_reader_lifecycle(ending: &str) {
use gateway_admin::ports::plugins::PluginStateLifecycle;
use std::time::Duration;
let Some(environment) = Environment::create().await else {
return;
};
let marker = environment.directory.path().join("usage-lifecycle");
let (runtime, core) = environment
.plugin(json!({"service":true,"mode":"usage_hold","records":marker}))
.await;
let admin = environment.bind_admin_accounts(&runtime, &core).await;
let services = admin.services();
runtime.bind_services(&services.public_services()).unwrap();
runtime.bind_usage_queries(&services.usage_data()).unwrap();
let registry = services.public_services();
let task = tokio::spawn(async move {
registry
.call("settings.admin_api_key_exists", json!(null))
.await
});
tokio::time::timeout(Duration::from_secs(5), async {
while !marker.with_extension("opened").exists() {
assert!(!task.is_finished(), "plugin must reach the open reader");
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
assert_eq!(environment.active_usage_transactions().await, 1);
match ending {
"cancel" => {
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
}
"disable" => {
let store = environment.store.admin_ports().plugins();
let snapshot = store.load_instances().await.unwrap();
services
.plugins()
.disable_instance(&snapshot.instances[0].id, &mutation())
.await
.unwrap();
assert!(!task.is_finished());
assert_eq!(environment.active_usage_transactions().await, 1);
// 停用阻止新调用进入插件,旧调用仍沿冻结计划完成
assert_eq!(
services
.public_services()
.call("settings.admin_api_key_exists", json!(null))
.await
.unwrap(),
json!(false)
);
std::fs::write(marker.with_extension("release"), b"release").unwrap();
assert_eq!(task.await.unwrap().unwrap(), json!(true));
}
"quiesce" => {
let snapshot = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap();
let instance = &snapshot.instances[0];
runtime
.quiesce_instance(&instance.id, &instance.artifact_sha256, instance.revision)
.await;
assert!(task.await.unwrap().is_err());
}
_ => unreachable!(),
}
tokio::time::timeout(Duration::from_secs(5), async {
while environment.active_usage_transactions().await != 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage transaction must be returned before runtime shutdown");
assert_eq!(
services
.usage_data()
.get_record(gateway_admin::model::usage_data::GetUsageRecord {
id: "held-usage".into(),
})
.await
.unwrap()
.id,
"held-usage"
);
drop(services);
drop(admin);
drop(core);
runtime.shutdown().await;
drop(runtime);
environment.close().await;
}
@@ -50,6 +50,50 @@ pub struct Environment {
}
impl Environment {
pub async fn active_usage_transactions(&self) -> i64 {
sqlx::query_scalar(
"select count(distinct a.pid) from pg_stat_activity a
join pg_locks l on l.pid=a.pid
join pg_class c on c.oid=l.relation
join pg_namespace n on n.oid=c.relnamespace
where n.nspname=$1 and c.relname='model_requests'
and a.xact_start is not null",
)
.bind(&self.schema)
.fetch_one(&self.admin)
.await
.unwrap()
}
pub async fn prepare_empty_analytics(&self) {
use gateway_store::postgres::{ObservabilityQueryBudget, PgRequestAnalyticsRepository};
use std::time::Duration;
let options = self
.admin
.connect_options()
.as_ref()
.clone()
.options([("search_path", self.schema.as_str())]);
let pool = PgPoolOptions::new()
.max_connections(1)
.connect_with(options)
.await
.unwrap();
let budget =
ObservabilityQueryBudget::try_new(1, Duration::from_secs(5), Duration::from_secs(10))
.unwrap();
let analytics = PgRequestAnalyticsRepository::new(pool.clone(), budget);
for _ in 0..10 {
analytics.maintain(100).await.unwrap();
if !analytics.status().await.unwrap().building {
pool.close().await;
return;
}
}
panic!("empty analytics did not become ready");
}
pub async fn client_key(&self, id: &str, plaintext: &str) {
self.client_key_with_limits(id, plaintext, gateway_core::policy::RateLimits::unlimited())
.await;
@@ -208,6 +208,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
));
}
use gateway_plugin_sdk::{call::services::settings, client::ServiceResponse};
if mode == "usage_hold" && call.is::<settings::ApiKeyExists>() {
return usage_data::hold(call, records.unwrap()).await;
}
if mode == "usage_data" && call.is::<settings::ApiKeyExists>() {
return usage_data::exercise(call).await;
}
@@ -6,32 +6,7 @@ use gateway_plugin_sdk::{
};
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 data = record_data();
let write = WriteUsageRecords {
changes: vec![UsageRecordChange::Put {
id: "sdk-usage".into(),
@@ -75,6 +50,27 @@ pub(super) async fn exercise(call: ServiceCall) -> Result<ServiceResponse, Plugi
.len(),
1
);
let scan_request = ScanUsageRecords {
filter: UsageReadFilter {
records: filter.clone(),
routing_group_ref: None,
},
batch_size: Some(1),
};
let mut reader = call.host.scan_usage_records(scan_request.clone()).await?;
let first = reader.next_batch().await?.unwrap();
let UsageReadItem::Record(fact) = &first.items[0] else {
panic!("record expected");
};
assert_eq!(fact.id, "sdk-usage");
assert_eq!(fact.data.client_key_ref, "sdk-key");
assert!(fact.usage_counted);
assert!(!first.complete);
assert!(reader.next_batch().await?.unwrap().complete);
assert!(reader.next_batch().await?.is_none());
let mut reader = call.host.scan_usage_records(scan_request).await?;
reader.close().await?;
assert!(reader.next_batch().await?.is_none());
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")
@@ -109,6 +105,33 @@ pub(super) async fn exercise(call: ServiceCall) -> Result<ServiceResponse, Plugi
})
.await
.unwrap();
let mut statistics = call
.host
.query_usage_statistics(QueryUsageStatistics {
filter: UsageReadFilter {
records: filter.clone(),
routing_group_ref: None,
},
group_by: UsageGrouping::Key,
batch_size: Some(1),
})
.await?;
let total = statistics.next_batch().await?.unwrap();
let UsageReadItem::Statistics(total) = &total.items[0] else {
panic!("statistics expected");
};
assert!(total.group.is_none());
assert_eq!(
total.metrics.values[&UsageMetric::RequestCount],
UsageCount(100)
);
let group = statistics.next_batch().await?.unwrap();
let UsageReadItem::Statistics(group) = &group.items[0] else {
panic!("statistics expected");
};
assert!(group.group.as_ref().unwrap().unassigned_contribution);
assert!(statistics.next_batch().await?.unwrap().complete);
assert!(statistics.next_batch().await?.is_none());
assert_eq!(
call.host
.service::<ListAggregates>(ListUsageAggregates {
@@ -144,3 +167,72 @@ pub(super) async fn exercise(call: ServiceCall) -> Result<ServiceResponse, Plugi
.unwrap();
ServiceResponse::from_result::<settings::ApiKeyExists>(Ok(true))
}
fn record_data() -> UsageRecordData {
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,
}
}
pub(super) async fn hold(
call: ServiceCall,
marker: std::path::PathBuf,
) -> Result<ServiceResponse, PluginFault> {
let call = call.into_typed::<settings::ApiKeyExists>()?;
call.host
.service::<WriteRecords>(WriteUsageRecords {
changes: vec![UsageRecordChange::Put {
id: "held-usage".into(),
data: Box::new(record_data()),
}],
})
.await
.unwrap();
let mut reader = call
.host
.scan_usage_records(ScanUsageRecords {
filter: UsageReadFilter {
records: UsageRecordFilter {
start_ms: 1780272000000,
end_ms: 1780272002000,
provider: None,
model: None,
account_ref: None,
client_key_ref: None,
},
routing_group_ref: None,
},
batch_size: Some(1),
})
.await?;
assert!(!reader.next_batch().await?.unwrap().complete);
std::fs::write(marker.with_extension("opened"), b"opened").unwrap();
while !marker.with_extension("release").exists() {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
// 故意不显式关闭,验证宿主在父调用终结时回收快照
ServiceResponse::from_result::<settings::ApiKeyExists>(Ok(true))
}
+3 -1
View File
@@ -19,7 +19,7 @@ SDK 包版本为 `0.1.0`,独立于网关版本,处于实验阶段。能力
| `call::frontend_authentication` | 数据面认证信封、认证器标识与外部 principal 结果 |
| `call::model` | 宿主模型调用的 canonical 事实、原生 wire 与有界事件编解码 |
| `call::middleware` | HTTP、WebSocket 消息、公开服务、模型请求、attempt 与上游出站准备的洋葱调用、单次 next 与惰性正文合同 |
| `call::services` | 公开服务的类型化操作、输入输出及错误合同;`usage` 提供[用量维护与计价](docs/usage-data.md) |
| `call::services` | 公开服务的类型化操作、输入输出及错误合同,`usage` 提供[用量查询、维护与计价](docs/usage-data.md) |
| `call::upstream_adapter` | 内置 OpenAI / xAI 的[受管上游适配器](docs/upstream-adapters.md)注册、执行、出站与续接合同 |
| `call::policy` | 模型路由、账号调度与重试决策合同 |
| `call::catalog` | 固定 Provider 的直接别名与适配模型注册合同,见[模型目录](docs/capabilities.md#模型目录) |
@@ -47,6 +47,8 @@ use gateway_plugin_sdk::client::{read_frame, write_frame};
`plugin.register`;作者不需要复制贡献项或手写方法名、阶段和敏感载荷编解码。
会话管理握手、有界并发、父调用关联、回调、流信用、取消、排空与关闭,容量与计时见[并发与期限](#并发与期限)
`HostClient::scan_usage_records` 与 `query_usage_statistics` 提供用量明细扫描和完整分组统计,按批读取与关闭见[用量查询](docs/usage-data.md#完整扫描与分组统计)
`ResponseStream::channel` 提供有界流生产入口,`ResponseStream::pull` 按消费进度驱动生产者;`PluginCall` 的宿主客户端只代表当前父调用,
不能把一个调用的资源误用到另一个调用。插件程序自行选择运行时和启动入口,不会启动网关。
`SessionConfig.maximum_stream_chunk_bytes` 默认 1 MiB,只约束单个业务流分块;
@@ -74,6 +74,11 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。宿主开放以下 11 类扩
所有回调都通过当前 `TypedCall.host` 发起,关联父调用及其取消信号。普通调用使用父调用剩余期限;已经建立的资源流按[会话生命周期](../README.md#帧与会话)回收,网络操作单独计时
用量扫描和完整分组统计使用 `HostClient::scan_usage_records`、`query_usage_statistics`,
对应 `host.usage.records.scan`、`host.usage.statistics.query`。返回的读取资源通过
`host.usage.read.next` 和 `host.usage.read.close` 拉取与关闭,SDK 提供 `UsageReader` 封装。
这些回调使用空控制参数和二进制 JSON 载荷,查询语义、快照与预算见[用量查询](usage-data.md#完整扫描与分组统计)
### 基础事实
基础事实接口提供账号分页、Key 分组绑定及额度观测。查询不进入模型执行链,也不触发计费;账号凭据使用[账号接口](#账号与凭据)读取
@@ -521,7 +526,7 @@ restore_and_validate_response(response).await
仅实现中间件时也可使用轻量的 `MiddlewarePlugin`;需要与管理或其他方法组合时使用
`PluginBuilder`,两者复用相同的类型化调用与结果合同
用量明细、聚合维护与费用试算使用 `call::services::usage`,输入输出和批量语义见 [用量维护与计价](usage-data.md)
用量查询、明细与聚合维护、费用试算使用 `call::services::usage`,输入输出和批量语义见 [用量查询、维护与计价](usage-data.md)
### 扩展的组合边界
@@ -1,8 +1,8 @@
# 用量维护与计价
# 用量查询、维护与计价
[返回 SDK](../README.md) · [公开服务](capabilities.md#公开服务)
`call::services::usage` 提供请求明细、聚合统计和 Token 费用试算。插件通过 `call.host.service::<Operation>(input)` 调用,沿用公开服务的调用身份、期限与取消约束
`call::services::usage` 提供请求明细、聚合统计和 Token 费用试算。单次操作通过 `call.host.service::<Operation>(input)` 调用,完整扫描与分组统计通过 `HostClient` 的分批读取方法调用,均绑定当前父调用
## 基础操作
@@ -15,7 +15,69 @@
| `WriteAggregates` | `usage.aggregates.write` | 按聚合键批量覆盖或删除统计贡献 |
| `EvaluatePricing` | `pricing.evaluate` | 根据当前价目或指定价格试算,逐项返回 `estimated` 或 `not_estimable` |
原生总量与趋势通过既有统计查询读取。明细列表只列请求记录,聚合贡献参与统计但不生成请求明细
明细列表只列请求记录,聚合贡献参与统计但不生成请求明细。完整分组与全局总量使用下方统计读取,趋势使用既有统计查询
## 完整扫描与分组统计
启用 `io` feature 后,`HostClient::scan_usage_records(ScanUsageRecords)` 与
`HostClient::query_usage_statistics(QueryUsageStatistics)` 返回 `UsageReader`,通过 `next_batch()` 逐批处理结果
```rust,no_run
use gateway_plugin_sdk::{PluginFault, client::HostClient};
use gateway_plugin_sdk::call::services::usage::{
QueryUsageStatistics, UsageGrouping, UsageReadBatch, UsageReadFilter,
};
async fn read_statistics(
host: &HostClient,
filter: UsageReadFilter,
mut consume: impl FnMut(UsageReadBatch),
) -> Result<(), PluginFault> {
let mut reader = host.query_usage_statistics(QueryUsageStatistics {
filter,
group_by: UsageGrouping::KeyModel,
batch_size: None,
}).await?;
while let Some(batch) = reader.next_batch().await? {
consume(batch);
}
Ok(())
}
```
### 范围与明细
- 时间按请求开始时间落入 UTC 毫秒 `[start_ms, end_ms)`,外部贡献按自身归属时间纳入
- 筛选支持 Provider、上游模型、账号、稳定 Client Key ID 及请求时保存的路由分组 ID。模型筛选与分组均使用实际 `model`,请求侧名称单独保存在 `requested_model`
- 明细按 `(started_at_ms, id)` 升序,包含 `UsageRecordData`、User-Agent、推理强度、Subagent kind、请求类型、已保存的路由范围与分组、恢复标记及统计资格
- `request_counted` 表示请求计数资格,`usage_counted` 表示完整交付用量资格,沿用宿主的预热和恢复关联规则
- 原始 Provider 观测与计价明细通过 `GetRecord` 获取,扫描只交付窄事实
### 分组与覆盖
`UsageGrouping` 支持 `model`、`key`、`key_model`、`routing_group`。首行 `group = null` 是全局总计,后续返回全部分组。`observation.total_groups` 表示分组数,排除总计行
路由分组使用记录中的 `routing_group_refs`,空列表进入未关联分组。多组请求分别计入各组,全局总计只计一次。分组筛选先选择匹配请求,再统计这些请求的全部分组归属
原生缺失身份使用 `null`。外部贡献缺少当前分组维度时,进入 `unassigned_contribution = true` 的独立分组并参与总计。显式维度筛选只匹配实际提供该维度的贡献,路由分组筛选排除外部贡献
- `values` 始终提供请求计数指标。Token 缺失表示未知,已知零保留为 `"0"`
- `native_request_count` 是原生请求数,`native_token_samples` 是各 Token 指标的原生已知请求数
- `contribution_count` 是外部贡献记录数,`contribution_token_samples` 是提供对应指标的贡献记录数,请求覆盖率由调用方结合来源含义解释
- `costs` 按原币种返回总额及原生、外部贡献金额。`cost_coverage` 返回原生完整交付请求的 Provider 上报、宿主计算和不可用费用数量
- 有符号贡献合并后若使任一分组计数或金额越界,读取返回错误,调用方据此判定整份结果未完成
### 快照、期限与关闭
一次读取使用固定只读快照,返回的 `generated_at_ms` 标记读取锚点。统计还返回投影覆盖、状态、滞后秒数和代次,结果反映快照内已发布投影,源记录的后续变化经 Worker 投影后可见
每批默认及最多 1024 行,编码后的 `UsageReadBatch` 最多 2 MiB。单行超过容量时返回 `capacity_exceeded`。总结果通过多批交付,只有收到 `complete = true` 才表示完整结束,SDK 随后返回 `None`
读取的总期限覆盖连接排队、SQL 和批次间消费等待,后续拉取共用同一期限。超时、取消、容量超限及存储失败返回错误,已经交付的批次属于未完成结果
`next_batch()` 的一次等待被取消后,再次调用会继续同一次在途读取。提前结束使用 `reader.close().await`,同一读取只允许一个消费者。句柄绑定插件实例及父调用,父调用结束或取消时宿主统一回收
普通停用阻止新调用进入实例,在途读取随原父调用完成。显式排空会取消在途调用。插件负责报表进度、重试策略和重新查询后的业务检查点
## 明细维护
@@ -37,7 +99,7 @@
聚合按 `bucket_start_ms` 这个归属时间纳入 `[start_ms, end_ms)` 查询及趋势时间桶,宿主不自动摊分。导入日汇总时,调用方选择该日的归属时间;需要小时分布时,调用方自行提供小时数据
查询只按实际提供的维度匹配,缺少的维度不会自动补全。未提供某个维度的贡献参与总量,但不进入该维度的分组或筛选;分组与总量是否对齐由调用方提供的数据决定。请求 ID、响应 ID、状态码、请求结果及文字搜索等明细筛选不包含聚合贡献;聚合也不生成延迟样本、尝试次数或请求列表页数。原生统计的无符号计数与非负金额若因贡献合计超出表示范围,查询返回错误,由调用方修正输入
查询只按实际提供的维度匹配,缺少的维度不会自动补全。未提供某个维度的贡献参与总量,完整分组读取使用独立未归属分组表达,其他维度查询只计入实际提供该维度的贡献。请求 ID、响应 ID、状态码、请求结果及文字搜索等明细筛选不包含聚合贡献,聚合也不生成延迟样本、尝试次数或请求列表页数。原生统计的无符号计数与非负金额若因贡献合计超出表示范围,查询返回错误,由调用方修正输入
聚合贡献参与原生总量、趋势、维度统计和账号用量,统计重建保留这些数据。贡献由调用方删除,不随请求历史保留任务清理;明细沿用请求历史保留策略
@@ -52,11 +114,11 @@
## 批量、错误与重试
- 每次写入或试算接受 1~500 项,编码后的输入不超过 2 MiB
- 查询每页默认 100 条,最多 200 条;明细页累计内容超过 2 MiB 时需减小页大小
- `ListRecords`、`ListAggregates` 每页默认 100 条,最多 200 条,明细页累计内容超过 2 MiB 时需减小页大小
- 时间为 UTC 毫秒,查询范围为左闭右开;游标不透明,仅用于相同筛选条件的后续页,分页不冻结跨调用快照
- `UsageCount`、`UsageDelta` 使用十进制字符串,避免 JavaScript 丢失整数精度;持久化数值受 PostgreSQL bigint 与费用精度范围限制
- 单次写入按输入顺序在一个事务中执行,任一项失败则整批回滚;跨批次进度由调用方保存
- 存储操作共享宿主观测连接与执行预算,超时、锁等待或取消会结束该次执行并释放资源
- 错误类型包括 `invalid_input`、`not_found`、`conflict`、`capacity_exceeded`、`unavailable`
- 错误类型包括 `invalid_input`、`not_found`、`conflict`、`capacity_exceeded`、`unavailable`,分批读取还区分 `cancelled` 与 `timeout`
响应丢失时,调用方按稳定 ID 或聚合键查询、重试覆盖。宿主不保存任务、来源登记、幂等回执、调整历史或撤销计划;调用方负责来源映射、业务去重、重试策略、并发写入协调和恢复。服务输入不能指定审计主体,宿主从当前调用身份生成普通变更审计
@@ -10,6 +10,8 @@ pub enum UsageErrorKind {
Conflict,
CapacityExceeded,
Unavailable,
Cancelled,
Timeout,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -20,3 +20,11 @@ pub use operations::*;
mod aggregate;
pub use aggregate::*;
mod read;
pub use read::*;
pub const RECORDS_SCAN: &str = "host.usage.records.scan";
pub const STATISTICS_QUERY: &str = "host.usage.statistics.query";
pub const READ_NEXT: &str = "host.usage.read.next";
pub const READ_CLOSE: &str = "host.usage.read.close";
@@ -0,0 +1,162 @@
//! 历史用量的窄事实、分组统计与调用内分批读取合同
use super::{UsageCount, UsageMetric, UsageRecordData, UsageRecordFilter};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
pub const USAGE_READ_MAX_BATCH_ROWS: u16 = 1024;
pub const USAGE_READ_MAX_BATCH_BYTES: usize = 2 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadFilter {
pub records: UsageRecordFilter,
pub routing_group_ref: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScanUsageRecords {
pub filter: UsageReadFilter,
pub batch_size: Option<u16>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageGrouping {
Model,
Key,
KeyModel,
RoutingGroup,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct QueryUsageStatistics {
pub filter: UsageReadFilter,
pub group_by: UsageGrouping,
pub batch_size: Option<u16>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum UsageReadQuery {
Records(ScanUsageRecords),
Statistics(QueryUsageStatistics),
}
impl UsageReadQuery {
pub fn filter(&self) -> &UsageReadFilter {
match self {
Self::Records(query) => &query.filter,
Self::Statistics(query) => &query.filter,
}
}
pub fn batch_size(&self) -> u16 {
match self {
Self::Records(query) => query.batch_size,
Self::Statistics(query) => query.batch_size,
}
.unwrap_or(USAGE_READ_MAX_BATCH_ROWS)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageRecordFact {
pub id: String,
pub data: UsageRecordData,
pub user_agent: Option<String>,
pub reasoning_effort: Option<String>,
pub subagent_kind: Option<String>,
pub request_kind: Option<String>,
pub routing_scope: Option<String>,
pub routing_group_refs: Vec<String>,
pub recovered: bool,
/// 由宿主现行规则派生,原始审计中的恢复请求仍可读取
pub request_counted: bool,
/// 包含完整交付、预热排除及恢复关联规则
pub usage_counted: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsGroup {
pub client_key_ref: Option<String>,
pub model: Option<String>,
pub routing_group_ref: Option<String>,
pub unassigned_contribution: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsCost {
pub currency: String,
pub amount: String,
pub native_amount: Option<String>,
/// 外部修正允许负值,原生费用来源覆盖单独返回
pub contribution_amount: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsCostCoverage {
pub provider_reported: UsageCount,
pub calculated: UsageCount,
pub unavailable: UsageCount,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsMetrics {
/// 缺失的 Token 指标表示未知,计数指标始终返回
pub values: BTreeMap<UsageMetric, UsageCount>,
pub native_request_count: UsageCount,
pub contribution_count: UsageCount,
pub native_token_samples: BTreeMap<UsageMetric, UsageCount>,
/// 数量表示提供该指标的贡献记录数,不能用作请求覆盖率的分母
pub contribution_token_samples: BTreeMap<UsageMetric, UsageCount>,
pub cost_coverage: UsageStatisticsCostCoverage,
pub costs: Vec<UsageStatisticsCost>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageStatisticsRow {
/// null 表示全局总计,分组重叠不会重复计入该行
pub group: Option<UsageStatisticsGroup>,
pub metrics: UsageStatisticsMetrics,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
pub enum UsageReadItem {
Record(Box<UsageRecordFact>),
Statistics(Box<UsageStatisticsRow>),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadAnalytics {
pub coverage_start_ms: Option<i64>,
pub coverage_end_ms: Option<i64>,
pub lag_seconds: u64,
pub state: String,
pub generation: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadObservation {
pub generated_at_ms: i64,
pub analytics: Option<UsageReadAnalytics>,
pub total_groups: Option<UsageCount>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadBatch {
pub observation: UsageReadObservation,
pub items: Vec<UsageReadItem>,
pub complete: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadOpened {
pub reader: Option<String>,
pub batch: UsageReadBatch,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UsageReadHandle {
pub reader: String,
}
@@ -10,6 +10,7 @@ mod plugin;
mod resources;
mod session;
mod upstream_adapter;
mod usage;
use crate::{ErrorCode, PluginFault};
use serde::{Serialize, de::DeserializeOwned};
@@ -33,6 +34,7 @@ pub use session::{
ResponseStream, SessionConfig, SessionError, StreamSender,
};
pub use upstream_adapter::{UpstreamWebSocket, UpstreamWebSocketUpgrade};
pub use usage::UsageReader;
async fn payload_call<T: Serialize, R: DeserializeOwned>(
host: &HostClient,
@@ -0,0 +1,97 @@
//! 用量查询的按需批次读取,取消等待保留同一次 RPC
use super::{HostClient, payload_call, pending_read::PendingRead};
use crate::{PluginFault, call::services::usage::*};
pub struct UsageReader {
host: HostClient,
reader: Option<String>,
first: Option<UsageReadBatch>,
read: PendingRead<Result<UsageReadBatch, PluginFault>>,
}
impl HostClient {
/// 在一次只读快照内扫描保留的原生请求事实
///
/// # Errors
/// 查询无效、存储超时或父调用终结时返回错误
pub async fn scan_usage_records(
&self,
request: ScanUsageRecords,
) -> Result<UsageReader, PluginFault> {
let opened = payload_call(self, RECORDS_SCAN, request).await?;
Ok(UsageReader::new(self.clone(), opened))
}
/// 查询完整分组统计,首行包含全局总计
///
/// # Errors
/// 查询无效、统计未就绪、存储超时或父调用终结时返回错误
pub async fn query_usage_statistics(
&self,
request: QueryUsageStatistics,
) -> Result<UsageReader, PluginFault> {
let opened = payload_call(self, STATISTICS_QUERY, request).await?;
Ok(UsageReader::new(self.clone(), opened))
}
}
impl UsageReader {
fn new(host: HostClient, opened: UsageReadOpened) -> Self {
Self {
host,
reader: opened.reader,
first: Some(opened.batch),
read: PendingRead::default(),
}
}
/// 按需读取下一批,完整结束后返回 None
/// 取消一次等待后再次调用会继续同一次读取,父调用结束时宿主回收句柄
///
/// # Errors
/// 超时、取消、容量超限和存储失败保留各自错误,不作为完整结束
pub async fn next_batch(&mut self) -> Result<Option<UsageReadBatch>, PluginFault> {
if let Some(first) = self.first.take() {
return Ok(Some(first));
}
let Some(reader) = &self.reader else {
return Ok(None);
};
let result = self
.read
.run(|| {
let host = self.host.clone();
let request = UsageReadHandle {
reader: reader.clone(),
};
async move { payload_call(&host, READ_NEXT, request).await }
})
.await;
if result.as_ref().map_or(true, |batch| batch.complete) {
self.reader = None;
}
result.map(Some)
}
/// 提前关闭当前读取,释放宿主快照与连接
///
/// # Errors
/// 父调用终结或句柄失效时返回错误
pub async fn close(&mut self) -> Result<(), PluginFault> {
if let Some(reader) = &self.reader {
payload_call::<_, ()>(
&self.host,
READ_CLOSE,
UsageReadHandle {
reader: reader.clone(),
},
)
.await?;
self.reader = None;
}
self.first = None;
self.read.clear();
Ok(())
}
}
@@ -5,3 +5,4 @@ mod http;
mod pending_read;
mod plugin;
mod session;
mod usage;
@@ -0,0 +1,141 @@
//! 用量读取的单次拉取、取消安全、显式失败及提前关闭
use super::session::{HostPeer, receive, send_call, shutdown, start_session};
use gateway_plugin_sdk::{
ErrorCode, Frame, Message, PluginFault,
call::services::usage::*,
client::{CallFuture, CallReply, PluginCall, PluginHandler, write_frame},
};
use serde_json::json;
struct UsageHandler;
impl PluginHandler for UsageHandler {
fn call(&self, call: PluginCall) -> CallFuture<'_> {
Box::pin(async move {
let mut reader = call
.host
.query_usage_statistics(QueryUsageStatistics {
filter: UsageReadFilter {
records: UsageRecordFilter {
start_ms: 0,
end_ms: 1000,
provider: None,
model: None,
account_ref: None,
client_key_ref: None,
},
routing_group_ref: None,
},
group_by: UsageGrouping::Model,
batch_size: Some(1),
})
.await?;
assert!(!reader.next_batch().await?.unwrap().complete);
{
let next = reader.next_batch();
tokio::pin!(next);
assert!(
std::future::poll_fn(|cx| std::task::Poll::Ready(next.as_mut().poll(cx)))
.await
.is_pending()
);
}
match call.method.as_str() {
"close" => {
reader.close().await?;
assert!(reader.next_batch().await?.is_none());
}
"error" => assert_eq!(
reader.next_batch().await.unwrap_err().code,
ErrorCode::Timeout
),
_ => {
assert!(reader.next_batch().await?.unwrap().complete);
assert!(reader.next_batch().await?.is_none());
}
}
Ok(CallReply::unary(json!(true), vec![]))
})
}
}
fn batch(complete: bool) -> UsageReadBatch {
UsageReadBatch {
observation: UsageReadObservation {
generated_at_ms: 10,
analytics: None,
total_groups: Some(UsageCount(1)),
},
items: vec![UsageReadItem::Statistics(Box::new(UsageStatisticsRow {
group: None,
metrics: UsageStatisticsMetrics::default(),
}))],
complete,
}
}
async fn callback(host: &mut HostPeer, expected: &str) -> u64 {
let Message::Callback { id, method, .. } = receive(host).await.message else {
panic!("callback expected");
};
assert_eq!(method, expected);
id
}
async fn reply(host: &mut HostPeer, id: u64, value: impl serde::Serialize) {
write_frame(
&mut host.writer,
&Frame {
message: Message::Result {
id,
result: json!({}),
},
payload: serde_json::to_vec(&value).unwrap(),
},
)
.await
.unwrap();
}
#[tokio::test]
async fn usage_reader_retains_pending_reads_and_distinguishes_failure_from_completion() {
for mode in ["complete", "error", "close"] {
let (mut host, task) = start_session(UsageHandler).await;
send_call(&mut host, 1, mode, json!({}), vec![]).await;
let id = callback(&mut host, STATISTICS_QUERY).await;
reply(
&mut host,
id,
UsageReadOpened {
reader: Some("reader".into()),
batch: batch(false),
},
)
.await;
let read_id = callback(&mut host, READ_NEXT).await;
match mode {
"close" => {
let close_id = callback(&mut host, READ_CLOSE).await;
reply(&mut host, close_id, ()).await;
}
"error" => write_frame(
&mut host.writer,
&Frame {
message: Message::Error {
id: read_id,
error: PluginFault::new(ErrorCode::Timeout, "read expired"),
},
payload: vec![],
},
)
.await
.unwrap(),
_ => reply(&mut host, read_id, batch(true)).await,
}
assert!(
matches!(receive(&mut host).await.message, Message::Result { id: 1, result } if result == json!(true))
);
shutdown(&mut host, task).await;
}
}
@@ -188,34 +188,13 @@ impl AnalyticsSnapshot<'_> {
.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());
apply_cost(state, &currency, &value)?;
}
}
Ok(())
}
}
fn apply_count(state: &mut MetricState, name: &str, delta: i128) -> StoreResult<()> {
pub(super) 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())
@@ -223,3 +202,25 @@ fn apply_count(state: &mut MetricState, name: &str, delta: i128) -> StoreResult<
state.counts.insert(name.into(), total);
Ok(())
}
pub(super) fn apply_cost(state: &mut MetricState, currency: &str, value: &str) -> StoreResult<()> {
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.to_owned(), amount.canonical());
Ok(())
}
@@ -11,7 +11,9 @@ mod incremental;
mod metrics;
mod query;
mod reset;
mod usage_query;
mod worker;
pub(crate) use usage_query::{declare_usage_statistics, decode_usage_statistics};
pub(crate) use query::*;
pub use worker::{AnalyticsMaintenance, AnalyticsStatus, PgRequestAnalyticsRepository};
@@ -33,7 +35,7 @@ pub(crate) async fn backup_projection_ready(
.map_err(|error| postgres_unavailable("check backup analytics readiness", error))
}
pub(super) async fn begin_snapshot(
pub(crate) async fn begin_snapshot(
connection: &mut ObservabilityConnection,
) -> StoreResult<sqlx::Transaction<'_, Postgres>> {
let query_timeout = connection.query_timeout();
@@ -0,0 +1,393 @@
//! 从互斥统计桶和窄事实生成完整分组游标,保留原生与外部贡献覆盖
use super::*;
use gateway_admin::model::usage_data::*;
use sqlx::postgres::PgRow;
const REQUEST_METRICS: &[(UsageMetric, &str)] = &[
(UsageMetric::RequestCount, "request_count"),
(UsageMetric::SuccessCount, "success_count"),
(UsageMetric::FailureCount, "failure_count"),
(UsageMetric::CancelledCount, "cancelled_count"),
(UsageMetric::IncompleteCount, "incomplete_count"),
(
UsageMetric::CompletedRequestCount,
"completed_request_count",
),
];
const TOKEN_FIELDS: &[(UsageMetric, &str)] = &[
(UsageMetric::InputTokens, "input_tokens"),
(UsageMetric::OutputTokens, "output_tokens"),
(UsageMetric::CachedInputTokens, "cached_input_tokens"),
(
UsageMetric::CacheWriteInputTokens,
"cache_write_input_tokens",
),
(
UsageMetric::ReasoningOutputTokens,
"reasoning_output_tokens",
),
(UsageMetric::ImageInputTokens, "image_input_tokens"),
(UsageMetric::ImageOutputTokens, "image_output_tokens"),
(UsageMetric::TotalTokens, "total_tokens"),
];
pub(crate) async fn declare_usage_statistics(
snapshot: &mut AnalyticsSnapshot<'_>,
request: &QueryUsageStatistics,
) -> StoreResult<()> {
let filter = &request.filter;
let range =
ObservabilityRange::new(time(filter.records.start_ms)?, time(filter.records.end_ms)?)?;
snapshot.require_visible(Some(range))?;
let collection = collection(request);
let mut query = QueryBuilder::<Postgres>::new(
"declare usage_read no scroll cursor for with native_parts as (select 'null'::text as name,'{}'::jsonb as counts,'{}'::jsonb as costs",
);
if let Some(collection) = collection {
let (hours, quarters) = snapshot.summary_buckets(range, true).await?;
for (table, seconds, buckets) in [
("request_analytics_1h", 3600, hours),
("request_analytics_15m", 900, quarters),
] {
if buckets.is_empty() {
continue;
}
let (key, model, provider, account) = summary_columns(collection);
let group = group_expression(request.group_by, key, model, "null::text", false);
query.push(format!(" union all select names.name,a.counts,a.costs from {table} a cross join lateral (values ({group}),('null')) names(name) where a.collection="))
.push_bind(collection).push(" and a.schema_version=").push_bind(SCHEMA_VERSION)
.push(" and a.histogram_version=").push_bind(HISTOGRAM_VERSION);
query::push_summary_buckets(&mut query, &buckets, seconds);
for (column, value) in [
(key, &filter.records.client_key_ref),
(model, &filter.records.model),
(provider, &filter.records.provider),
(account, &filter.records.account_ref),
] {
if let Some(value) = value {
query
.push(format!(" and {column}="))
.push_bind(value.clone());
}
}
}
}
let spans = if collection.is_some() {
snapshot.uncovered_spans(range).await?
} else {
vec![(range.start, range.end)]
};
if !spans.is_empty() {
query.push(" union all select name,counts,costs from (with matching as not materialized (select mr.*,coalesce((");
query.push(completed_usage_fact_predicate("mr"));
query
.push("),false) as is_completed_usage,")
.push_bind(range.start)
.push("::timestamptz as output_bucket,names.name from request_analytics_facts mr");
push_fact_groups(&mut query, request.group_by);
query.push(" where (");
for (index, (start, end)) in spans.into_iter().enumerate() {
if index > 0 {
query.push(" or ");
}
query
.push("(mr.started_at >= ")
.push_bind(start)
.push(" and mr.started_at < ")
.push_bind(end)
.push(")");
}
query.push(")");
push_unrecovered_request_filter(&mut query, "mr");
super::super::super::usage_facts::push_usage_read_filter(&mut query, filter);
query.push(")");
let counters = counter_names();
let names = counters.iter().map(String::as_str).collect::<Vec<_>>();
aggregate::push_fact_aggregation_selected(
&mut query,
false,
true,
DistributionRead::None,
Some(&names),
None,
);
query.push(") facts");
}
query.push("), inputs as (select name,'native'::text as origin,counts,costs from native_parts");
if filter.routing_group_ref.is_none() {
let group = group_expression(
request.group_by,
"nullif(a.dimensions[4],'')",
"nullif(a.dimensions[2],'')",
"null::text",
true,
);
query.push(format!(" union all select names.name,'contribution',a.counts,a.costs from request_analytics_1h a cross join lateral (values ({group}),('null')) names(name) where a.collection='usage_contribution' and a.bucket_start >= "))
.push_bind(range.start).push(" and a.bucket_start < ").push_bind(range.end);
for (index, value) in [
(1, &filter.records.provider),
(2, &filter.records.model),
(3, &filter.records.account_ref),
(4, &filter.records.client_key_ref),
] {
if let Some(value) = value {
query
.push(format!(" and a.dimensions[{index}]="))
.push_bind(value.clone());
}
}
}
// 分组和全局总计在同一查询中生成;每条输入只贡献一次全局计数
// 汇总在数据库完成,应用只接收当前批次,金额按原币种保留
query.push(r#"), groups as (
select name, count(*) filter(where origin='contribution')::bigint as contribution_count
from inputs group by name
), metric_values as (
select name, origin, m.key, sum(m.value::numeric)::text as value, count(*)::bigint as samples
from inputs cross join lateral jsonb_each_text(counts) m
group by name, origin, m.key
), metric_groups as (
select name,
jsonb_object_agg(key,value) filter(where origin='native') as native_counts,
jsonb_object_agg(key,value) filter(where origin='contribution') as contribution_counts,
jsonb_object_agg(key,samples) filter(where origin='contribution') as contribution_samples
from metric_values group by name
), cost_values as (
select name, origin, c.key, trim_scale(sum(c.value::numeric))::text as value
from inputs cross join lateral jsonb_each_text(costs) c
where c.key not in ('standard_usd','no_cache_usd')
group by name, origin, c.key
), cost_groups as (
select name,
jsonb_object_agg(key,value) filter(where origin='native') as native_costs,
jsonb_object_agg(key,value) filter(where origin='contribution') as contribution_costs
from cost_values group by name
) select g.name, g.contribution_count, count(*) over()-1 as total_groups,
coalesce(m.native_counts,'{}'::jsonb) as native_counts,
coalesce(m.contribution_counts,'{}'::jsonb) as contribution_counts,
coalesce(m.contribution_samples,'{}'::jsonb) as contribution_samples,
coalesce(c.native_costs,'{}'::jsonb) as native_costs,
coalesce(c.contribution_costs,'{}'::jsonb) as contribution_costs
from groups g
left join metric_groups m using(name)
left join cost_groups c using(name)
order by (g.name<>'null'), g.name collate "C""#);
query
.build()
.execute(&mut *snapshot.tx)
.await
.map_err(|error| postgres_unavailable("open usage statistics cursor", error))?;
Ok(())
}
fn collection(query: &QueryUsageStatistics) -> Option<&'static str> {
let filter = &query.filter;
if query.group_by == UsageGrouping::RoutingGroup || filter.routing_group_ref.is_some() {
return None;
}
let key = matches!(query.group_by, UsageGrouping::Key | UsageGrouping::KeyModel)
|| filter.records.client_key_ref.is_some();
if filter.records.account_ref.is_some() {
return (!key).then_some("account_model");
}
if key {
if query.group_by == UsageGrouping::Key
&& filter.records.model.is_none()
&& filter.records.provider.is_none()
{
Some("key")
} else {
Some("key_model")
}
} else {
Some("provider_model")
}
}
fn summary_columns(collection: &str) -> (&'static str, &'static str, &'static str, &'static str) {
match collection {
"key" => (
"nullif(a.dimensions[1],'')",
"null::text",
"null::text",
"null::text",
),
"key_model" => (
"nullif(a.dimensions[1],'')",
"nullif(a.dimensions[4],'')",
"nullif(a.dimensions[2],'')",
"null::text",
),
"account_model" => (
"null::text",
"nullif(a.dimensions[4],'')",
"nullif(a.dimensions[2],'')",
"nullif(a.dimensions[1],'')",
),
_ => (
"null::text",
"nullif(a.dimensions[3],'')",
"nullif(a.dimensions[1],'')",
"null::text",
),
}
}
fn group_expression(
grouping: UsageGrouping,
key: &str,
model: &str,
group: &str,
contribution: bool,
) -> String {
let (key, model, group) = match grouping {
UsageGrouping::Model => ("null::text", model, "null::text"),
UsageGrouping::Key => (key, "null::text", "null::text"),
UsageGrouping::KeyModel => (key, model, "null::text"),
UsageGrouping::RoutingGroup => ("null::text", "null::text", group),
};
let unassigned = if contribution {
match grouping {
UsageGrouping::Model => format!("({model}) is null"),
UsageGrouping::Key => format!("({key}) is null"),
UsageGrouping::KeyModel => format!("({key}) is null or ({model}) is null"),
UsageGrouping::RoutingGroup => "true".to_owned(),
}
} else {
"false".to_owned()
};
format!(
"jsonb_build_object('client_key_ref',{key},'model',{model},'routing_group_ref',{group},'unassigned_contribution',({unassigned}))::text"
)
}
fn push_fact_groups(query: &mut QueryBuilder<Postgres>, grouping: UsageGrouping) {
let group = group_expression(
grouping,
"mr.client_api_key_ref",
"mr.upstream_model_id",
"r.group_ref",
false,
);
if grouping == UsageGrouping::RoutingGroup {
query.push(format!(" cross join lateral (select {group} as name from unnest(case when cardinality(mr.routing_group_refs)>0 then mr.routing_group_refs else array[null::text] end) r(group_ref) union all select 'null') names"));
} else {
query.push(format!(
" cross join lateral (values ({group}),('null')) names(name)"
));
}
}
fn counter_names() -> Vec<String> {
REQUEST_METRICS
.iter()
.chain(TOKEN_FIELDS)
.map(|(_, name)| (*name).to_owned())
.chain(
TOKEN_FIELDS
.iter()
.map(|(_, name)| format!("{name}_known_count")),
)
.chain(
[
"provider_reported_count",
"calculated_count",
"unavailable_count",
]
.map(str::to_owned),
)
.collect()
}
fn time(milliseconds: i64) -> StoreResult<DateTime<Utc>> {
DateTime::from_timestamp_millis(milliseconds).ok_or_else(|| invalid("invalid usage query time"))
}
pub(crate) fn decode_usage_statistics(row: PgRow) -> StoreResult<UsageStatisticsRow> {
let name: String = get(&row, "name")?;
let native_counts: BTreeMap<String, String> =
get::<sqlx::types::Json<_>>(&row, "native_counts")?.0;
let contribution_counts: BTreeMap<String, String> =
get::<sqlx::types::Json<_>>(&row, "contribution_counts")?.0;
let contribution_samples: BTreeMap<String, u64> =
get::<sqlx::types::Json<_>>(&row, "contribution_samples")?.0;
let native_costs: BTreeMap<String, String> =
get::<sqlx::types::Json<_>>(&row, "native_costs")?.0;
let contribution_costs: BTreeMap<String, String> =
get::<sqlx::types::Json<_>>(&row, "contribution_costs")?.0;
let mut state = MetricState::default();
for (key, value) in native_counts {
state.counts.insert(
key,
value
.parse()
.map_err(|_| invalid("native usage count exceeds result range"))?,
);
}
let native_request_count = UsageCount(state.count("request_count"));
let cost_coverage = UsageStatisticsCostCoverage {
provider_reported: UsageCount(state.count("provider_reported_count")),
calculated: UsageCount(state.count("calculated_count")),
unavailable: UsageCount(state.count("unavailable_count")),
};
let native_token_samples = TOKEN_FIELDS
.iter()
.map(|(metric, name)| {
(
*metric,
UsageCount(state.count(&format!("{name}_known_count"))),
)
})
.collect::<BTreeMap<_, _>>();
for (name, value) in contribution_counts {
contributions::apply_count(
&mut state,
&name,
value
.parse()
.map_err(|_| invalid("aggregate counter exceeds arithmetic range"))?,
)?;
}
state.costs.clone_from(&native_costs);
for (currency, amount) in &contribution_costs {
contributions::apply_cost(&mut state, currency, amount)?;
}
let mut values = REQUEST_METRICS
.iter()
.map(|(metric, name)| (*metric, UsageCount(state.count(name))))
.collect::<BTreeMap<_, _>>();
let mut contribution_token_samples = BTreeMap::new();
for (metric, name) in TOKEN_FIELDS {
let samples = contribution_samples.get(*name).copied().unwrap_or_default();
contribution_token_samples.insert(*metric, UsageCount(samples));
if native_token_samples[metric].get() > 0 || samples > 0 {
values.insert(*metric, UsageCount(state.count(name)));
}
}
let costs = state
.costs()?
.into_iter()
.map(|cost| UsageStatisticsCost {
native_amount: native_costs.get(&cost.currency).cloned(),
contribution_amount: contribution_costs.get(&cost.currency).cloned(),
currency: cost.currency,
amount: cost.amount.as_str().to_owned(),
})
.collect();
Ok(UsageStatisticsRow {
group: serde_json::from_str(&name).map_err(|_| invalid("invalid usage group identity"))?,
metrics: UsageStatisticsMetrics {
values,
native_request_count,
contribution_count: UsageCount(
u64::try_from(get::<i64>(&row, "contribution_count")?)
.map_err(|_| invalid("negative contribution count"))?,
),
native_token_samples,
contribution_token_samples,
costs,
cost_coverage,
},
})
}
@@ -37,6 +37,7 @@ pub(crate) use analytics::accounts::{
pub(crate) use analytics::backup_projection_ready;
pub(crate) use analytics::{AGGREGATE_COLLECTION, AnalyticsSnapshot};
pub use analytics::{AnalyticsMaintenance, AnalyticsStatus, PgRequestAnalyticsRepository};
pub(crate) use analytics::{begin_snapshot, declare_usage_statistics, decode_usage_statistics};
pub(crate) use mapping::*;
pub use model::*;
pub(crate) use queries::*;
@@ -1,7 +1,9 @@
//! 复用请求明细和聚合存储的有界维护入口,不管理插件工作流
mod aggregates;
mod read;
mod records;
mod scan;
use super::{ObservabilityConnection, ObservabilityQueryBudget};
use crate::{StoreError, StoreResult, postgres_unavailable};
@@ -23,6 +25,13 @@ impl PgUsageDataStore {
#[async_trait::async_trait]
impl UsageDataStore for PgUsageDataStore {
async fn open_read(
&self,
query: UsageReadQuery,
cancellation: gateway_core::lifecycle::CancellationToken,
) -> UsageDataResult<Box<dyn UsageReadSession>> {
Ok(read::open(self, query, cancellation))
}
async fn get_record(&self, request: GetUsageRecord) -> UsageDataResult<UsageRecord> {
self.budget
.run(&self.pool, "usage record", |mut connection| async move {
@@ -0,0 +1,253 @@
//! 在观测总预算内持有只读游标,按消费需求交付有界批次
use super::*;
use futures::TryStreamExt;
use gateway_core::lifecycle::CancellationToken;
use sqlx::postgres::PgRow;
use tokio::sync::{mpsc, oneshot};
type BatchReply = oneshot::Sender<UsageReadBatch>;
struct Reader {
requests: mpsc::Sender<BatchReply>,
task: Option<tokio::task::JoinHandle<UsageDataResult<()>>>,
cancellation: CancellationToken,
}
pub(super) fn open(
store: &PgUsageDataStore,
query: UsageReadQuery,
cancellation: CancellationToken,
) -> Box<dyn UsageReadSession> {
let (requests, receiver) = mpsc::channel(1);
let store = store.clone();
let cancellation = cancellation.child_token();
let worker_cancellation = cancellation.clone();
let task = tokio::spawn(async move {
let run = store
.budget
.run(&store.pool, "usage read", |mut connection| async move {
drive(&mut connection, query, receiver).await
});
tokio::select! {
biased;
() = worker_cancellation.cancelled() => Err(UsageDataError::new(UsageErrorKind::Cancelled, "用量读取已取消")),
result = run => result.map_err(read_error),
}
});
Box::new(Reader {
requests,
task: Some(task),
cancellation,
})
}
#[async_trait::async_trait]
impl UsageReadSession for Reader {
async fn next_batch(&mut self) -> UsageDataResult<UsageReadBatch> {
let (reply, result) = oneshot::channel();
if self.requests.send(reply).await.is_ok()
&& let Ok(batch) = result.await
{
return Ok(batch);
}
let Some(task) = self.task.take() else {
return Err(UsageDataError::new(
UsageErrorKind::NotFound,
"用量读取已结束",
));
};
task.await
.map_err(|_| UsageDataError::unavailable("用量读取任务已中断"))??;
Err(UsageDataError::new(
UsageErrorKind::NotFound,
"用量读取已结束",
))
}
fn is_finished(&self) -> bool {
self.task
.as_ref()
.is_none_or(tokio::task::JoinHandle::is_finished)
}
async fn close(&mut self) {
self.cancellation.cancel();
if let Some(task) = self.task.take() {
let _ = task.await;
}
}
}
impl Drop for Reader {
fn drop(&mut self) {
self.cancellation.cancel();
}
}
async fn open_cursor<'a>(
connection: &'a mut ObservabilityConnection,
query: &UsageReadQuery,
) -> StoreResult<(Transaction<'a, Postgres>, UsageReadObservation)> {
let (mut tx, observation) = match query {
UsageReadQuery::Records(request) => {
let mut tx = super::super::observability::begin_snapshot(connection).await?;
scan::declare(&mut tx, request).await?;
let generated_at: DateTime<Utc> = sqlx::query_scalar("select transaction_timestamp()")
.fetch_one(&mut *tx)
.await
.map_err(db_error)?;
(
tx,
UsageReadObservation {
generated_at_ms: generated_at.timestamp_millis(),
analytics: None,
total_groups: None,
},
)
}
UsageReadQuery::Statistics(request) => {
let mut snapshot =
super::super::observability::AnalyticsSnapshot::new(connection).await?;
super::super::observability::declare_usage_statistics(&mut snapshot, request).await?;
let metadata = snapshot.metadata();
let observation = UsageReadObservation {
generated_at_ms: metadata.generated_at.timestamp_millis(),
analytics: Some(UsageReadAnalytics {
coverage_start_ms: metadata.coverage_start.map(|time| time.timestamp_millis()),
coverage_end_ms: metadata.coverage_end.map(|time| time.timestamp_millis()),
lag_seconds: metadata.lag_seconds,
state: metadata.state.as_str().to_owned(),
generation: metadata.generation,
}),
total_groups: None,
};
(snapshot.tx, observation)
}
};
sqlx::query("set transaction read only")
.execute(&mut *tx)
.await
.map_err(db_error)?;
Ok((tx, observation))
}
async fn drive(
connection: &mut ObservabilityConnection,
query: UsageReadQuery,
mut requests: mpsc::Receiver<BatchReply>,
) -> StoreResult<()> {
let (mut tx, mut observation) = open_cursor(connection, &query).await?;
// 游标行流始终由本任务持有,等待下一批不会重建查询或丢弃半途读取
let (reply, batch) = {
// 游标名称复用而结果列随查询改变,FETCH 每次取得当前列描述
let mut rows = sqlx::query("fetch all from usage_read")
.persistent(false)
.fetch(&mut *tx);
let mut pending = None;
loop {
let Some(reply) = requests.recv().await else {
return Ok(());
};
let mut batch = UsageReadBatch {
observation: observation.clone(),
items: Vec::new(),
complete: false,
};
let mut bytes = 0;
while batch.items.len() < usize::from(query.batch_size()) {
let item = match pending.take() {
Some(item) => item,
None => {
let Some(row) = rows.try_next().await.map_err(db_error)? else {
batch.complete = true;
break;
};
decode(row, &query, &mut observation)?
}
};
let size = serde_json::to_vec(&item)
.map_err(|_| invalid("encode usage read item"))?
.len()
+ 1;
if batch.items.is_empty() {
batch.observation = observation.clone();
bytes = serde_json::to_vec(&batch)
.map_err(|_| invalid("encode usage batch"))?
.len();
}
if bytes + size > USAGE_READ_MAX_BATCH_BYTES {
if batch.items.is_empty() {
return Err(invalid("usage read item exceeds batch capacity")
.with_source(ReadCapacity));
}
pending = Some(item);
break;
}
bytes += size;
batch.items.push(item);
}
batch.observation = observation.clone();
if batch.complete {
break (reply, batch);
}
if reply.send(batch).is_err() {
return Ok(());
}
}
};
// 完整结束只在数据库事务完成后发布,提交失败不能表现为正常 EOF
tx.commit().await.map_err(db_error)?;
let _ = reply.send(batch);
Ok(())
}
#[derive(Debug, thiserror::Error)]
#[error("usage read item exceeds batch capacity")]
struct ReadCapacity;
fn read_error(error: StoreError) -> UsageDataError {
if let StoreError::InvalidData {
source: Some(source),
..
} = &error
&& source.is::<ReadCapacity>()
{
return UsageDataError::new(UsageErrorKind::CapacityExceeded, "单条用量结果超过批次容量");
}
if let StoreError::Unavailable {
source: Some(source),
..
} = &error
{
let sql_timeout = source
.downcast_ref::<sqlx::Error>()
.and_then(sqlx::Error::as_database_error)
.and_then(|error| error.code())
.is_some_and(|code| code == "57014");
if source.is::<tokio::time::error::Elapsed>() || sql_timeout {
return UsageDataError::new(UsageErrorKind::Timeout, "用量读取超过执行期限");
}
}
public_error(error)
}
fn decode(
row: PgRow,
query: &UsageReadQuery,
observation: &mut UsageReadObservation,
) -> StoreResult<UsageReadItem> {
match query {
UsageReadQuery::Records(_) => {
scan::decode(row).map(|record| UsageReadItem::Record(Box::new(record)))
}
UsageReadQuery::Statistics(_) => {
let groups: i64 = row.try_get("total_groups").map_err(db_error)?;
observation.total_groups = Some(UsageCount(
u64::try_from(groups).map_err(|_| invalid("negative usage group count"))?,
));
super::super::observability::decode_usage_statistics(row)
.map(|statistics| UsageReadItem::Statistics(Box::new(statistics)))
}
}
}
@@ -92,6 +92,15 @@ pub(super) async fn list(
Ok(UsageRecordPage { items, next_cursor })
}
fn decode(row: PgRow) -> StoreResult<UsageRecord> {
Ok(UsageRecord {
id: row.try_get("id").map_err(db_error)?,
data: decode_data(&row)?,
billing_snapshot: row.try_get("billing_snapshot_json").map_err(db_error)?,
provider_observation: row.try_get("provider_observation_json").map_err(db_error)?,
})
}
pub(super) fn decode_data(row: &PgRow) -> StoreResult<UsageRecordData> {
let token = |name: &str| -> StoreResult<Option<UsageCount>> {
row.try_get::<Option<i64>, _>(name)
.map_err(db_error)?
@@ -128,44 +137,40 @@ fn decode(row: PgRow) -> StoreResult<UsageRecord> {
},
_ => 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,
Ok(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")?,
},
billing_snapshot: row.try_get("billing_snapshot_json").map_err(db_error)?,
provider_observation: row.try_get("provider_observation_json").map_err(db_error)?,
cost,
})
}
fn cost_parts(cost: &UsageCost) -> (&str, Option<&str>, Option<&str>) {
match cost {
UsageCost::Unavailable => ("unavailable", None, None),
@@ -0,0 +1,57 @@
//! 统计扫描的原生请求窄投影,沿用宿主用量资格与稳定身份
use super::*;
use sqlx::{QueryBuilder, postgres::PgRow};
pub(super) async fn declare(
tx: &mut Transaction<'_, Postgres>,
request: &ScanUsageRecords,
) -> StoreResult<()> {
let mut query = QueryBuilder::<Postgres>::new(format!(
"declare usage_read no scroll cursor for
select mr.id, mr.client_api_key_ref, mr.provider_kind, mr.provider_account_ref,
mr.requested_model_id, mr.upstream_model_id, mr.protocol, mr.endpoint,
mr.operation, mr.client_transport, mr.service_tier, mr.started_at, mr.completed_at,
mr.outcome, mr.input_tokens, mr.output_tokens, mr.cached_input_tokens,
mr.cache_write_input_tokens, mr.reasoning_output_tokens, mr.image_input_tokens,
mr.image_output_tokens, mr.total_tokens, mr.cost_source, mr.cost_currency,
mr.cost_amount::text as cost_amount, mr.user_agent, mr.reasoning_effort,
mr.subagent_kind, mr.request_kind, mr.routing_scope, mr.routing_group_refs,
mr.recovered_at is not null as recovered,
coalesce(({}),false) and mr.recovered_at is null as usage_counted
from model_request_observations mr where mr.started_at >= ",
super::super::completed_usage_fact_predicate("mr"),
));
query
.push_bind(at(request.filter.records.start_ms)?)
.push(" and mr.started_at < ")
.push_bind(at(request.filter.records.end_ms)?);
super::super::usage_facts::push_usage_read_filter(&mut query, &request.filter);
query.push(" order by mr.started_at,mr.id");
query.build().execute(&mut **tx).await.map_err(db_error)?;
Ok(())
}
pub(super) fn decode(row: PgRow) -> StoreResult<UsageRecordFact> {
let user_agent = row.try_get("user_agent").map_err(db_error)?;
let reasoning_effort = row.try_get("reasoning_effort").map_err(db_error)?;
let subagent_kind = row.try_get("subagent_kind").map_err(db_error)?;
let request_kind = row.try_get("request_kind").map_err(db_error)?;
let routing_scope = row.try_get("routing_scope").map_err(db_error)?;
let routing_group_refs = row.try_get("routing_group_refs").map_err(db_error)?;
let recovered = row.try_get("recovered").map_err(db_error)?;
let usage_counted = row.try_get("usage_counted").map_err(db_error)?;
Ok(UsageRecordFact {
id: row.try_get("id").map_err(db_error)?,
data: records::decode_data(&row)?,
user_agent,
reasoning_effort,
subagent_kind,
request_kind,
routing_scope,
routing_group_refs,
recovered,
request_counted: !recovered,
usage_counted,
})
}
@@ -1,5 +1,6 @@
//! `model_requests` 完成用量事实的 PostgreSQL 语义
use gateway_admin::model::usage_data::UsageReadFilter;
use sqlx::{Postgres, QueryBuilder};
/// 只把已完整交付给客户端、且存在用量证据的成功响应投影为用量事实
@@ -55,3 +56,24 @@ pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres
pub(crate) fn push_unrecovered_request_filter(query: &mut QueryBuilder<Postgres>, alias: &str) {
query.push(format!(" and {alias}.recovered_at is null"));
}
pub(crate) fn push_usage_read_filter(query: &mut QueryBuilder<Postgres>, filter: &UsageReadFilter) {
for (column, value) in [
("mr.provider_kind", &filter.records.provider),
("mr.upstream_model_id", &filter.records.model),
("mr.provider_account_ref", &filter.records.account_ref),
("mr.client_api_key_ref", &filter.records.client_key_ref),
] {
if let Some(value) = value {
query
.push(format!(" and {column}="))
.push_bind(value.clone());
}
}
if let Some(group) = &filter.routing_group_ref {
query
.push(" and ")
.push_bind(group.clone())
.push("=any(mr.routing_group_refs)");
}
}
@@ -1,5 +1,7 @@
//! 通用用量维护的原生统计、重试、原子性、重建与资源预算验收
mod read;
use super::{TestDatabase, drain_analytics, observability_query_budget, observability_repository};
use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::{
@@ -0,0 +1,480 @@
//! 用量读取的完整分批、统计覆盖、一致快照与取消预算验证
use super::*;
use gateway_core::lifecycle::CancellationToken;
use std::{collections::BTreeSet, time::Instant};
fn read_filter() -> UsageReadFilter {
UsageReadFilter {
records: filter(),
routing_group_ref: None,
}
}
fn scan(batch_size: u16) -> UsageReadQuery {
UsageReadQuery::Records(ScanUsageRecords {
filter: read_filter(),
batch_size: Some(batch_size),
})
}
async fn statistics(
store: &PgUsageDataStore,
filter: UsageReadFilter,
group_by: UsageGrouping,
) -> Vec<UsageStatisticsRow> {
let mut reader = store
.open_read(
UsageReadQuery::Statistics(QueryUsageStatistics {
filter,
group_by,
batch_size: Some(17),
}),
CancellationToken::new(),
)
.await
.unwrap();
let mut rows = Vec::new();
let mut observation = None;
loop {
let batch = reader.next_batch().await.unwrap();
assert!(serde_json::to_vec(&batch).unwrap().len() <= USAGE_READ_MAX_BATCH_BYTES);
if let Some(observation) = &observation {
assert_eq!(observation, &batch.observation);
}
observation = Some(batch.observation.clone());
for item in batch.items {
let UsageReadItem::Statistics(row) = item else {
panic!("statistics row expected");
};
rows.push(*row);
}
if batch.complete {
break;
}
}
assert_eq!(
observation.unwrap().total_groups.unwrap().get(),
(rows.len() - 1) as u64
);
assert!(rows[0].group.is_none());
rows
}
#[tokio::test]
async fn scan_delivers_ten_thousand_records_from_one_snapshot() {
let Some(db) = TestDatabase::create("usage_scan").await else {
return;
};
let store = PgUsageDataStore::new(
db.pool.clone(),
ObservabilityQueryBudget::try_new(2, Duration::from_secs(5), Duration::from_secs(60))
.unwrap(),
);
for base in (0..10_000).step_by(500) {
store
.write_records(records(
(base..base + 500)
.map(|index| put(&format!("request_{index:05}"), 10))
.collect(),
))
.await
.unwrap();
}
sqlx::query("update model_requests set request_observation_json=jsonb_set(request_observation_json,'{request}',(request_observation_json->'request') || '{\"userAgent\":\"usage-test\",\"reasoningEffort\":\"high\",\"subagentKind\":\"worker\"}'::jsonb)").execute(&db.pool).await.unwrap();
let baseline_started = Instant::now();
let mut cursor = None;
let mut baseline_batches = 0;
loop {
let page = store
.list_records(ListUsageRecords {
filter: filter(),
cursor,
limit: Some(200),
})
.await
.unwrap();
baseline_batches += 1;
cursor = page.next_cursor;
if cursor.is_none() {
break;
}
}
let baseline_elapsed = baseline_started.elapsed();
let started = Instant::now();
let mut reader = store
.open_read(scan(1024), CancellationToken::new())
.await
.unwrap();
let first = reader.next_batch().await.unwrap();
let observation = first.observation.clone();
store
.write_records(records(vec![
UsageRecordChange::Delete {
id: "request_09999".into(),
},
put("request_new", 99),
put("request_09000", 99),
]))
.await
.unwrap();
let mut batch = first;
let mut ids = BTreeSet::new();
let mut batches = 0;
loop {
assert_eq!(batch.observation, observation);
batches += 1;
for item in batch.items {
let UsageReadItem::Record(record) = item else {
panic!("record expected");
};
assert_eq!(record.data.tokens.total_tokens, Some(UsageCount(10)));
assert_eq!(record.user_agent.as_deref(), Some("usage-test"));
assert_eq!(record.reasoning_effort.as_deref(), Some("high"));
assert_eq!(record.subagent_kind.as_deref(), Some("worker"));
assert!(record.request_counted && record.usage_counted);
assert!(ids.insert(record.id));
}
if batch.complete {
break;
}
batch = reader.next_batch().await.unwrap();
}
assert_eq!(ids.len(), 10_000);
assert!(ids.contains("request_09999") && !ids.contains("request_new"));
eprintln!(
"usage scan rows=10000 baseline_batches={baseline_batches} baseline_ms={} scan_batches={batches} scan_ms={}",
baseline_elapsed.as_millis(),
started.elapsed().as_millis()
);
let current = store
.get_record(GetUsageRecord {
id: "request_09000".into(),
})
.await
.unwrap();
assert_eq!(current.data.tokens.total_tokens, Some(UsageCount(99)));
reader.close().await;
db.close().await;
}
#[tokio::test]
async fn statistics_preserve_all_groups_unknowns_contributions_and_overlapping_routes() {
let Some(db) = TestDatabase::create("usage_statistics").await else {
return;
};
let store = PgUsageDataStore::new(db.pool.clone(), observability_query_budget());
let mut changes = Vec::new();
for index in 0..137 {
let mut value = data(10);
value.client_key_ref = format!("key_{index}");
value.requested_model = Some("alias".into());
value.model = Some(format!("model_{}", index % 2));
value.started_at_ms += 1000;
value.completed_at_ms = Some(value.started_at_ms + 1000);
value.tokens.cached_input_tokens = Some(UsageCount(0));
value.cost = UsageCost::Known {
source: UsageCostSource::Calculated,
amount: "1".into(),
currency: if index % 2 == 0 { "USD" } else { "EUR" }.into(),
};
changes.push(UsageRecordChange::Put {
id: format!("row_{index}"),
data: Box::new(value),
});
}
let mut prewarm = data(999);
prewarm.client_key_ref = "key_0".into();
prewarm.model = Some("model_0".into());
prewarm.started_at_ms += 1000;
prewarm.completed_at_ms = Some(prewarm.started_at_ms + 1000);
changes.push(UsageRecordChange::Put {
id: "prewarm".into(),
data: Box::new(prewarm),
});
store.write_records(records(changes)).await.unwrap();
sqlx::query("update model_requests set request_kind='prewarm' where id='prewarm'")
.execute(&db.pool)
.await
.unwrap();
sqlx::query("update model_requests set request_observation_json=jsonb_set(request_observation_json,'{routing}','{\"scope\":\"groups\",\"groupRefs\":[\"g1\",\"g2\"],\"groupNamesSnapshot\":[\"one\",\"two\"]}'::jsonb) where id='row_0'").execute(&db.pool).await.unwrap();
let UsageAggregateChange::Put { mut aggregate } = contribution("known", 2, 20, "0.5") else {
unreachable!();
};
aggregate.identity.bucket_start_ms += 1000;
aggregate.identity.dimensions.client_key_ref = Some("key_0".into());
aggregate.identity.dimensions.model = Some("model_0".into());
let mut unknown = aggregate.clone();
unknown.identity.key = "unknown".into();
unknown.identity.dimensions = UsageDimensions::default();
unknown.counts = BTreeMap::from([(UsageMetric::RequestCount, UsageDelta(3))]);
unknown.costs.clear();
store
.write_aggregates(aggregates(vec![
UsageAggregateChange::Put { aggregate },
UsageAggregateChange::Put { aggregate: unknown },
]))
.await
.unwrap();
drain_analytics(&db.pool).await;
let mut filter = read_filter();
filter.records.start_ms += 500;
filter.records.end_ms -= 500;
let rows = statistics(&store, filter.clone(), UsageGrouping::KeyModel).await;
let projected = statistics(&store, read_filter(), UsageGrouping::KeyModel).await;
assert_eq!(rows.len(), projected.len());
for (edge, full) in rows.iter().zip(&projected) {
assert_eq!(edge, full);
}
assert_eq!(
statistics(&store, read_filter(), UsageGrouping::Key)
.await
.len(),
139
);
assert_eq!(
statistics(&store, read_filter(), UsageGrouping::Model)
.await
.len(),
4
);
let mut selected = read_filter();
selected.records.client_key_ref = Some("key_0".into());
let selected_rows = statistics(&store, selected, UsageGrouping::Model).await;
assert_eq!(
selected_rows[0].metrics.values[&UsageMetric::RequestCount].get(),
4
);
assert_eq!(
selected_rows[0].metrics.values[&UsageMetric::TotalTokens].get(),
30
);
assert_eq!(rows.len(), 139);
let total = &rows[0].metrics;
assert_eq!(total.values[&UsageMetric::RequestCount].get(), 143);
assert_eq!(total.values[&UsageMetric::CompletedRequestCount].get(), 139);
assert_eq!(total.values[&UsageMetric::TotalTokens].get(), 1390);
assert_eq!(
total.native_token_samples[&UsageMetric::InputTokens].get(),
137
);
assert_eq!(
total.contribution_token_samples[&UsageMetric::TotalTokens].get(),
1
);
assert!(!total.values.contains_key(&UsageMetric::OutputTokens));
assert_eq!(total.values[&UsageMetric::CachedInputTokens].get(), 0);
assert_eq!(total.cost_coverage.calculated.get(), 137);
assert_eq!(
total
.costs
.iter()
.find(|cost| cost.currency == "USD")
.unwrap()
.amount,
"69.5"
);
assert_eq!(
total
.costs
.iter()
.find(|cost| cost.currency == "EUR")
.unwrap()
.amount,
"68"
);
let routes = statistics(&store, filter.clone(), UsageGrouping::RoutingGroup).await;
assert_eq!(routes[0].metrics, *total);
for group in ["g1", "g2"] {
let row = routes
.iter()
.find(|row| {
row.group
.as_ref()
.is_some_and(|key| key.routing_group_ref.as_deref() == Some(group))
})
.unwrap();
assert_eq!(row.metrics.values[&UsageMetric::RequestCount].get(), 1);
}
filter.routing_group_ref = Some("g1".into());
let scoped = statistics(&store, filter.clone(), UsageGrouping::Key).await;
assert_eq!(
scoped[0].metrics.values[&UsageMetric::RequestCount].get(),
1
);
assert_eq!(scoped[0].metrics.contribution_count.get(), 0);
filter.routing_group_ref = None;
filter.records.model = Some("alias".into());
let empty = statistics(&store, filter, UsageGrouping::Model).await;
assert_eq!(empty.len(), 1);
assert_eq!(empty[0].metrics.values[&UsageMetric::RequestCount].get(), 0);
let UsageAggregateChange::Put { mut aggregate } = contribution("negative-group", -100, 0, "0")
else {
unreachable!();
};
aggregate.identity.dimensions.client_key_ref = Some("zzz".into());
aggregate.identity.dimensions.model = Some("model_0".into());
aggregate.counts = BTreeMap::from([(UsageMetric::RequestCount, UsageDelta(-100))]);
aggregate.costs.clear();
store
.write_aggregates(aggregates(vec![UsageAggregateChange::Put { aggregate }]))
.await
.unwrap();
let mut reader = store
.open_read(
UsageReadQuery::Statistics(QueryUsageStatistics {
filter: read_filter(),
group_by: UsageGrouping::KeyModel,
batch_size: Some(1),
}),
CancellationToken::new(),
)
.await
.unwrap();
assert!(!reader.next_batch().await.unwrap().complete);
let error = loop {
match reader.next_batch().await {
Ok(batch) => assert!(!batch.complete, "negative group must fail the read"),
Err(error) => break error,
}
};
assert_eq!(error.kind, UsageErrorKind::InvalidInput);
reader.close().await;
// 合法贡献各自有界,合并后的单个统计行仍须服从交付字节上限
let currencies = (b'A'..=b'Z')
.flat_map(|first| {
(b'A'..=b'Z').flat_map(move |second| {
(b'A'..=b'Z')
.map(move |third| String::from_utf8(vec![first, second, third]).unwrap())
})
})
.collect::<Vec<_>>();
let native_changes = currencies
.iter()
.take(800)
.enumerate()
.map(|(index, currency)| {
let mut value = data(1);
value.model = Some("large-costs".into());
value.cost = UsageCost::Known {
source: UsageCostSource::Calculated,
amount: "0.1234567891".into(),
currency: currency.clone(),
};
UsageRecordChange::Put {
id: format!("cost-{index}"),
data: Box::new(value),
}
})
.collect::<Vec<_>>();
for changes in native_changes.chunks(500) {
store
.write_records(records(changes.to_vec()))
.await
.unwrap();
}
drain_analytics(&db.pool).await;
let changes = currencies
.chunks(32)
.enumerate()
.map(|(index, currencies)| UsageAggregateChange::Put {
aggregate: UsageAggregate {
identity: UsageAggregateKey {
bucket_start_ms: read_filter().records.start_ms,
dimensions: UsageDimensions {
model: Some("large-costs".into()),
..Default::default()
},
key: format!("costs-{index}"),
},
counts: BTreeMap::new(),
costs: currencies
.iter()
.map(|currency| (currency.clone(), "1234567890.1234567891".into()))
.collect(),
},
})
.collect::<Vec<_>>();
for changes in changes.chunks(500) {
store
.write_aggregates(aggregates(changes.to_vec()))
.await
.unwrap();
}
let mut large_filter = read_filter();
large_filter.records.model = Some("large-costs".into());
let mut reader = store
.open_read(
UsageReadQuery::Statistics(QueryUsageStatistics {
filter: large_filter,
group_by: UsageGrouping::Model,
batch_size: Some(1),
}),
CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(
reader.next_batch().await.unwrap_err().kind,
UsageErrorKind::CapacityExceeded
);
reader.close().await;
db.close().await;
}
#[tokio::test]
async fn byte_boundaries_timeout_and_cancellation_release_readers() {
let Some(db) = TestDatabase::create("usage_read_budget").await else {
return;
};
let budget =
ObservabilityQueryBudget::try_new(1, Duration::from_secs(2), Duration::from_millis(500))
.unwrap();
let store = PgUsageDataStore::new(db.pool.clone(), budget);
store
.write_records(records(
(0..8)
.map(|index| put(&format!("large_{index}"), 1))
.collect(),
))
.await
.unwrap();
sqlx::query("update model_requests set request_observation_json=jsonb_set(request_observation_json,'{request,userAgent}',to_jsonb(repeat('x',700000)))").execute(&db.pool).await.unwrap();
let mut reader = store
.open_read(scan(1024), CancellationToken::new())
.await
.unwrap();
let first = reader.next_batch().await.unwrap();
assert!(!first.complete && first.items.len() < 8);
assert!(serde_json::to_vec(&first).unwrap().len() <= USAGE_READ_MAX_BATCH_BYTES);
reader.close().await;
let cancellation = CancellationToken::new();
let mut reader = store
.open_read(scan(1), cancellation.clone())
.await
.unwrap();
assert_eq!(reader.next_batch().await.unwrap().items.len(), 1);
cancellation.cancel();
assert_eq!(
reader.next_batch().await.unwrap_err().kind,
UsageErrorKind::Cancelled
);
let mut reader = store
.open_read(scan(1), CancellationToken::new())
.await
.unwrap();
reader.next_batch().await.unwrap();
tokio::time::sleep(Duration::from_millis(600)).await;
assert_eq!(
reader.next_batch().await.unwrap_err().kind,
UsageErrorKind::Timeout
);
let mut after = store
.open_read(scan(1), CancellationToken::new())
.await
.unwrap();
assert_eq!(after.next_batch().await.unwrap().items.len(), 1);
after.close().await;
db.close().await;
}
+8 -3
View File
@@ -236,7 +236,7 @@ flowchart LR
- **调用预算由阶段合同统一解释**:SDK 定义场景默认值,插件在 `Ready` 声明所需预算,Runtime 与 SDK 解析并冻结到会话;各入口按实例和阶段读取,宿主不提供部署侧业务时长上限。
注册在进程启动与握手完成后使用自己的预算;CLI 的执行与结果提交共用一次期限,重试委托链共享匹配实例中最长的预算并受请求剩余期限约束,单次调用仍受自身声明限制。
业务预算不延长进程启动与握手、排队、取消确认、关闭和外部请求期限;声明与默认值见 [SDK 调用预算](../backend/crates/gateway-plugin/sdk/README.md#声明调用预算)
- **资源跟随调用**:账号、HTTP、模型、状态和日志回调绑定有效父调用,取消随调用传播。HTTP 与模型 request/attempt 中间件在首个流式响应头交付后,由正文 owner 管理取消与终态,不沿用初始处理期限;网络操作与模型执行各自计时。
- **资源跟随调用**:账号、HTTP、模型、用量读取、状态和日志回调绑定有效父调用,取消随调用传播。HTTP 与模型 request/attempt 中间件在首个流式响应头交付后,由正文 owner 管理取消与终态,不沿用初始处理期限;网络操作与模型执行各自计时。
Runtime 与 SDK 以父处理任务的实际挂起和恢复授权管理执行名额,不按回调方法决定是否让出;
`next` 仅在父任务挂起且没有其他应计时回调时暂停处理预算,暂停覆盖恢复排队。
执行调度由 Runtime 统一负责,准入和恢复等待使用独立期限;初次排队、总在途上下文和驻留回调各有上限,`next` 与普通回调共享驻留额度。
@@ -894,11 +894,16 @@ Admin 的手动重置复用同一账本与 Key 行锁,在一个事务中清零
### 用量维护服务
Admin 定义明细、聚合维护与纯计价合同,Provider 复用原生 Token 计价公式,Runtime 传递可信调用身份。插件负责外部格式、来源映射、去重、调整历史、目标选择和跨调用流程
Admin 定义明细扫描、完整分组统计、用量维护与纯计价合同,Provider 复用原生 Token 计价公式,Runtime 传递可信调用身份。插件负责报表、外部格式、来源映射、去重、调整历史、目标选择和跨调用流程
明细直接维护 `model_requests` 的终态记录,通过统计失效触发器更新投影;计价快照在维护后清空,Key 结算账本保持独立。聚合贡献使用 `request_analytics_1h` 的 `usage_contribution` 集合保存调用方输入,其他集合为请求投影。桶重建和版本重建保留贡献集合,查询在同一快照中按指定维度和归属时间合并贡献;聚合的删除由调用方管理
Store 在单次事务内完成有限批量读写和普通审计,复用观测连接与执行预算。原生统计结果遵守各自数值类型,越界返回错误,不自动修正插件数据。调用合同与时间归属见 [SDK 用量维护与计价](../backend/crates/gateway-plugin/sdk/docs/usage-data.md)
Store 在单次事务内完成有限批量读写和普通审计,复用观测连接与执行预算。原生统计结果遵守各自数值类型,越界返回错误,不自动修正插件数据。调用合同与时间归属见 [SDK 用量查询、维护与计价](../backend/crates/gateway-plugin/sdk/docs/usage-data.md)
完整读取在同一只读快照中持有数据库游标,Store 按需交付有界批次,预算覆盖 SQL 和消费等待。
分组统计复用互斥汇总桶与边界窄事实,路由分组使用请求时保存的分组 ID,总计独立去重。
Runtime 在父调用资源表中持有读取句柄并传播取消,SDK 保留中断等待时的在途读取。
正常完成在事务收尾后确认,提前关闭、错误和父调用结束均释放读取资源
### 模型价格与费用快照