mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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:
@@ -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);
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
+24
-23
@@ -188,34 +188,13 @@ impl AnalyticsSnapshot<'_> {
|
||||
.map(|v| v.0)
|
||||
.unwrap_or_default();
|
||||
for (currency, value) in costs {
|
||||
let current = state
|
||||
.costs
|
||||
.get(¤cy)
|
||||
.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, ¤cy, &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;
|
||||
}
|
||||
@@ -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 保留中断等待时的在途读取。
|
||||
正常完成在事务收尾后确认,提前关闭、错误和父调用结束均释放读取资源
|
||||
|
||||
### 模型价格与费用快照
|
||||
|
||||
|
||||
Reference in New Issue
Block a user