starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-07 00:05:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): restore account usage and request observability

Browse Source 
Restore provider-native quota and reasoning observations, complete request lifecycle accounting, and align account and usage surfaces with the restored semantics.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-25 00:33:12 +08:00
co-authored by Codex
parent 24c7c575d2
commit 75e90d41b6
41 changed files with 2305 additions and 407 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-admin/src/model/accounts.rs
+19 -1
View File
@@ -7,7 +7,7 @@ use futures::Stream;
use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision};
use super::{PageSize, Revision, observability::TimeRange};
/// Provider 账号当前可用性。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -151,6 +151,24 @@ pub struct AccountUsage {
pub models: Vec<AccountModelUsage>,
}
/// 某个账号在调用方指定时间窗口内的本地用量查询。
///
/// `key` 只用于把聚合结果关联回调用方的窗口，不承载 Provider 私有语义。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUsageWindowQuery {
pub account_id: String,
pub key: String,
pub range: TimeRange,
}
/// 一个账号时间窗口用量查询的结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUsageWindowResult {
pub account_id: String,
pub key: String,
pub usage: AccountUsage,
}
/// 账号列表页所需的完整存储事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPage {
backend/crates/gateway-admin/src/model/observability.rs
+17
View File
@@ -327,6 +327,21 @@ pub struct ProviderBillingInput {
pub total: CurrencyCost,
}
/// 已完整交付且由 Provider 计算费用的持久请求事实。
///
/// 控制面仅保留通用事实，具体 Provider 负责校验已持久化总额并恢复标准费用。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageCalculatedBillingFact {
pub bucket_start: DateTime<Utc>,
pub provider_kind: String,
pub upstream_model_id: String,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
pub cache_write_tokens: Option<u64>,
pub total: CurrencyCost,
}
/// Provider 已确认的逐项费用与单价。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CalculatedBillingBreakdown {
@@ -549,6 +564,8 @@ pub struct UsageRecord {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
/// Provider 已筛选的专有观测 JSON；管理领域保持不透明。
pub provider_metadata_json: Option<String>,
pub attempt_count: u32,
pub upstream_send_state: String,
pub downstream_committed_at: Option<DateTime<Utc>>,
backend/crates/gateway-admin/src/ports/store.rs
+16 -4
View File
@@ -10,8 +10,8 @@ use chrono::{DateTime, Utc};
use crate::model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -20,8 +20,8 @@ use crate::model::{
},
observability::{
DashboardObservation, DashboardRuntimeSlots, DiagnosticDimension, DiagnosticObservation,
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter,
UsageOverview, UsagePage, UsageQuery,
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageCalculatedBillingFact,
UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
},
provider_credentials::{
AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult,
@@ -90,6 +90,11 @@ pub trait AccountStore: Send + Sync {
account_ids: &[String],
) -> AdminStoreResult<Vec<AccountUsage>>;
async fn load_account_usage_by_windows(
&self,
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>>;
async fn list_credentials(
&self,
provider_kind: &gateway_core::routing::ProviderKind,
@@ -249,6 +254,13 @@ pub trait ObservabilityStore: Send + Sync {
filter: UsageFilter,
) -> AdminStoreResult<Vec<RequestMetricPoint>>;
/// 返回可由 Provider 重新校验的已计算费用事实，用于恢复标准费用趋势。
async fn usage_calculated_billing_facts(
&self,
range: TimeRange,
filter: UsageFilter,
) -> AdminStoreResult<Vec<UsageCalculatedBillingFact>>;
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage>;
async fn usage_record_detail(&self, request_id: &str) -> AdminStoreResult<UsageDetail>;
backend/crates/gateway-admin/src/use_case/accounts.rs
+70 -4
View File
@@ -19,12 +19,13 @@ use crate::{
AdminError, MutationContext, Revision,
accounts::{
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
AccountRecord, AccountStatus,
AccountRecord, AccountStatus, AccountUsageWindowQuery,
},
observability::TimeRange,
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
PrepareCredentialRefresh, ProviderModels, ProviderQuotaRequest,
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
ProviderQuotaWindow,
},
},
ports::{
@@ -132,6 +133,45 @@ impl DefaultAccountsService {
Ok((account, provider))
}
async fn attach_quota_local_usage(
&self,
accounts: &[AccountRecord],
quotas: &mut [ProviderQuota],
) -> Result<(), AdminError> {
let windows = accounts
.iter()
.zip(quotas.iter())
.flat_map(|(account, quota)| {
quota
.windows
.iter()
.filter_map(|window| quota_usage_window(&account.id, window))
})
.collect::<Vec<_>>();
if windows.is_empty() {
return Ok(());
}
let usage_by_window = self
.accounts
.load_account_usage_by_windows(&windows)
.await
.map_err(|error| map_store_error(error, "quota window usage"))?
.into_iter()
.map(|result| ((result.account_id, result.key), result.usage))
.collect::<BTreeMap<_, _>>();
for (account, quota) in accounts.iter().zip(quotas) {
for window in &mut quota.windows {
if window.local_usage.is_none() {
let key = (account.id.clone(), window.key.clone());
if let Some(usage) = usage_by_window.get(&key) {
window.local_usage = Some(usage.clone());
}
}
}
}
Ok(())
}
async fn load_directory_item(
&self,
account_id: &ProviderAccountId,
@@ -165,7 +205,7 @@ impl DefaultAccountsService {
},
)?;
let rolling_usage = rolling_usage.into_iter().next();
let quota = provider
let mut quota = provider
.quota(ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: refresh_quota,
@@ -173,6 +213,11 @@ impl DefaultAccountsService {
})
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
self.attach_quota_local_usage(
std::slice::from_ref(&account),
std::slice::from_mut(&mut quota),
)
.await?;
Ok(AccountDirectoryItem {
status: account_status(&account, now),
usage: usage.into_iter().next(),
@@ -214,7 +259,7 @@ impl AccountsService for DefaultAccountsService {
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let quotas = futures::future::join_all(page.items.iter().map(|account| async {
let mut quotas = futures::future::join_all(page.items.iter().map(|account| async {
let account_id = ProviderAccountId::new(account.id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
let provider = self
@@ -233,6 +278,8 @@ impl AccountsService for DefaultAccountsService {
.await
.into_iter()
.collect::<Result<Vec<_>, AdminError>>()?;
self.attach_quota_local_usage(&page.items, &mut quotas)
.await?;
let items = page
.items
.into_iter()
@@ -451,3 +498,22 @@ fn retained_usage_range(now: chrono::DateTime<Utc>, retention_days: u32) -> Time
end: now,
}
}
fn quota_usage_window(
account_id: &str,
window: &ProviderQuotaWindow,
) -> Option<AccountUsageWindowQuery> {
if window.local_usage.is_some() {
return None;
}
let reset_at = window.reset_at?;
let seconds = i64::try_from(window.window_seconds?).ok()?;
let start = reset_at.checked_sub_signed(Duration::seconds(seconds))?;
let range = TimeRange::new(start, reset_at).ok()?;
// 上游百分比以该 reset 边界定义；以当前时间回推会让本地 Token 属于另一窗口。
Some(AccountUsageWindowQuery {
account_id: account_id.to_owned(),
key: window.key.clone(),
range,
})
}
backend/crates/gateway-admin/src/use_case/observability.rs
+92 -11
View File
@@ -1,6 +1,9 @@
//! 观测查询、趋势、费用覆盖与健康阈值规则。
use std::{collections::BTreeMap, sync::Arc};
use std::{
collections::{BTreeMap, btree_map::Entry},
sync::Arc,
};
use async_trait::async_trait;
use chrono::{DateTime, Duration, Timelike as _, Utc};
@@ -14,15 +17,16 @@ use crate::{
DashboardPeriodMetrics, DashboardResult, DecimalAmount, DiagnosticDimension,
DiagnosticsItem, DiagnosticsResult, HealthStatus, HealthTimeline, HealthTimelinePoint,
OpsErrorPage, OpsErrorQuery, ProviderBillingInput, RequestMetricPoint, RequestMetrics,
TimeRange, Trend, TrendKind, TrendPoint, TrendSummary, UsageBilling, UsageDetail,
UsageFilter, UsageInsights, UsageInsightsCost, UsageInsightsCostPoint,
UsageInsightsHealth, UsageInsightsHealthPoint, UsageInsightsPerformance,
UsageInsightsPerformancePoint, UsageOverview, UsagePage, UsageQuery, UsageSummary,
TimeRange, Trend, TrendKind, TrendPoint, TrendSummary, UsageBilling,
UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageInsights, UsageInsightsCost,
UsageInsightsCostPoint, UsageInsightsHealth, UsageInsightsHealthPoint,
UsageInsightsPerformance, UsageInsightsPerformancePoint, UsageOverview, UsagePage,
UsageQuery, UsageSummary,
},
provider_credentials::ProviderQuotaRequest,
},
ports::{
provider::ProviderAdminRegistry,
provider::{ProviderAdminErrorKind, ProviderAdminRegistry},
store::{ObservabilityStore, SettingsStore},
},
};
@@ -215,12 +219,14 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
filter: UsageFilter,
) -> Result<UsageInsights, AdminError> {
let (overview, trend) = futures::try_join!(
let (overview, trend, billing_facts) = futures::try_join!(
self.store.usage_summary(range, filter.clone()),
self.store.usage_trend(range, filter),
self.store.usage_trend(range, filter.clone()),
self.store.usage_calculated_billing_facts(range, filter),
)
.map_err(|error| map_store_error(error, "usage insights"))?;
build_usage_insights(overview, trend)
let standard_costs = recover_standard_costs(&self.providers, billing_facts)?;
build_usage_insights(overview, trend, standard_costs)
}
async fn diagnostics(
@@ -268,6 +274,7 @@ impl ObservabilityService for DefaultObservabilityService {
fn build_usage_insights(
overview: UsageOverview,
trend: Vec<RequestMetricPoint>,
standard_costs: UsageStandardCosts,
) -> Result<UsageInsights, AdminError> {
let granularity = trend.first().map_or(
crate::model::observability::Granularity::FifteenMinutes,
@@ -325,12 +332,16 @@ fn build_usage_insights(
.collect(),
};
let estimated_cost = usd_cost(&overview.attempts.costs);
let UsageStandardCosts {
total: standard_cost,
by_bucket: standard_costs_by_bucket,
} = standard_costs;
let cost = UsageInsightsCost {
cost_per_request: estimated_cost
.as_ref()
.and_then(|cost| cost.checked_div_u64(requests.request_count)),
estimated_cost,
standard_cost: None,
standard_cost,
tokens_per_request: rate_or_zero(requests.total_tokens, requests.request_count),
cached_token_rate: rate_or_zero(requests.cached_tokens, requests.input_tokens),
cache_hit_request_rate: ratio(
@@ -350,7 +361,7 @@ fn build_usage_insights(
cached_tokens: point.metrics.cached_tokens,
total_tokens: point.metrics.total_tokens,
estimated_cost: usd_cost(&point.costs),
standard_cost: None,
standard_cost: standard_costs_by_bucket.get(&point.bucket_start).cloned(),
cached_token_rate: rate_or_zero(
point.metrics.cached_tokens,
point.metrics.input_tokens,
@@ -374,6 +385,76 @@ fn build_usage_insights(
})
}
#[derive(Debug, Default)]
struct UsageStandardCosts {
total: Option<DecimalAmount>,
by_bucket: BTreeMap<DateTime<Utc>, DecimalAmount>,
}
impl UsageStandardCosts {
fn add(
&mut self,
bucket_start: DateTime<Utc>,
amount: DecimalAmount,
) -> Result<(), AdminError> {
self.total = Some(match self.total.take() {
Some(current) => current
.checked_add(&amount)
.ok_or_else(|| AdminError::internal("Usage standard cost exceeds decimal range"))?,
None => amount.clone(),
});
match self.by_bucket.entry(bucket_start) {
Entry::Occupied(mut entry) => {
let sum = entry.get().checked_add(&amount).ok_or_else(|| {
AdminError::internal("Usage standard cost exceeds decimal range")
})?;
entry.insert(sum);
}
Entry::Vacant(entry) => {
entry.insert(amount);
}
}
Ok(())
}
}
fn recover_standard_costs(
providers: &ProviderAdminRegistry,
facts: Vec<UsageCalculatedBillingFact>,
) -> Result<UsageStandardCosts, AdminError> {
let mut standard_costs = UsageStandardCosts::default();
for fact in facts {
let Ok(provider_kind) = ProviderKind::new(fact.provider_kind) else {
continue;
};
let input = ProviderBillingInput {
upstream_model_id: fact.upstream_model_id,
input_tokens: fact.input_tokens,
output_tokens: fact.output_tokens,
cached_tokens: fact.cached_tokens,
cache_write_tokens: fact.cache_write_tokens,
total: fact.total,
};
let breakdown = match providers.calculated_billing(&provider_kind, &input) {
Ok(breakdown) => breakdown,
Err(error) if error.kind() == ProviderAdminErrorKind::Unsupported => continue,
Err(error) => return Err(map_provider_error(error, "usage billing")),
};
let Some(breakdown) = breakdown else {
continue;
};
if !breakdown
.standard_amount
.currency
.eq_ignore_ascii_case("USD")
{
continue;
}
standard_costs.add(fact.bucket_start, breakdown.standard_amount.amount)?;
}
Ok(standard_costs)
}
fn usd_cost(costs: &[CurrencyCost]) -> Option<DecimalAmount> {
costs
.iter()
backend/crates/gateway-admin/tests/use_case/accounts.rs
+107 -8
View File
@@ -18,7 +18,8 @@ use gateway_admin::{
MutationContext, Revision,
accounts::{
AccountAvailability, AccountListQuery, AccountPage, AccountRecord, AccountSummary,
AccountUsage, DeleteAccounts, SetAccountEnabled,
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, DeleteAccounts,
SetAccountEnabled,
},
observability::TimeRange,
provider_credentials::{
@@ -31,7 +32,7 @@ use gateway_admin::{
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
ProviderQuotaRequest,
ProviderQuotaRequest, ProviderQuotaWindow,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -59,6 +60,7 @@ pub(super) struct FakeProviderAdmin {
export_inputs: Mutex<Vec<ProviderExportCredentialInput>>,
import_account_ids: Mutex<Vec<String>>,
quota_requests: Mutex<Vec<ProviderQuotaRequest>>,
quota: Mutex<ProviderQuota>,
}
impl FakeProviderAdmin {
@@ -72,6 +74,7 @@ impl FakeProviderAdmin {
export_inputs: Mutex::new(Vec::new()),
import_account_ids: Mutex::new(vec!["acct_prepared".to_owned()]),
quota_requests: Mutex::new(Vec::new()),
quota: Mutex::new(empty_quota()),
})
}
@@ -100,6 +103,10 @@ impl FakeProviderAdmin {
.clone()
}
pub(super) fn set_quota(&self, quota: ProviderQuota) {
*self.quota.lock().expect("provider quota") = quota;
}
pub(super) fn pending(&self) -> Option<PendingAuthorizationMutation> {
self.pending.lock().expect("pending authorization").clone()
}
@@ -282,12 +289,7 @@ impl ProviderAdmin for FakeProviderAdmin {
{
return Err(ProviderAdminError::new(kind));
}
Ok(ProviderQuota {
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
provider_data: None,
})
Ok(self.quota.lock().expect("provider quota").clone())
}
async fn models(
@@ -327,6 +329,7 @@ pub(super) struct FakeAccountStore {
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
authorization_revision: Mutex<Option<Revision>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
}
impl FakeAccountStore {
@@ -337,6 +340,7 @@ impl FakeAccountStore {
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
authorization_revision: Mutex::new(None),
quota_window_usage: Mutex::new(Vec::new()),
})
}
@@ -355,6 +359,10 @@ impl FakeAccountStore {
.expect("authorization revision")
}
pub(super) fn set_quota_window_usage(&self, usage: Vec<AccountUsageWindowResult>) {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
fn record(&self, event: &'static str) {
self.events.lock().expect("store events").push(event);
}
@@ -405,6 +413,17 @@ impl AccountStore for FakeAccountStore {
Ok(Vec::new())
}
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
Ok(self
.quota_window_usage
.lock()
.expect("quota window usage")
.clone())
}
async fn list_credentials(
&self,
provider_kind: &ProviderKind,
@@ -758,6 +777,54 @@ async fn accounts_list_should_return_complete_directory_semantics() {
);
}
#[tokio::test]
async fn accounts_list_should_attach_local_usage_to_quota_windows() {
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::hours(1);
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(reset_at),
local_usage: None,
provider_data: None,
}],
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
store.set_quota_window_usage(vec![AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: "primary".to_owned(),
usage: quota_local_usage("acct_test", 4_330_000),
}]);
let page = accounts_service(provider, store)
.await
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
search: None,
status: None,
sort: None,
})
.await
.expect("quota window usage");
let usage = page.items[0].quota.windows[0]
.local_usage
.as_ref()
.expect("quota window local usage");
assert_eq!(usage.total_tokens, Some(4_330_000));
}
#[tokio::test]
async fn accounts_refresh_provider_failure_should_not_call_store_commit() {
let events = events();
@@ -832,6 +899,38 @@ pub(super) fn document() -> ProviderDocument {
ProviderDocument::new(OpaqueProviderData::new(Default::default()))
}
fn empty_quota() -> ProviderQuota {
ProviderQuota {
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
provider_data: None,
}
}
fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage {
AccountUsage {
account_id: account_id.to_owned(),
request_count: 1,
success_count: 1,
input_tokens: Some(total_tokens),
output_tokens: Some(0),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
image_request_count: 0,
image_request_failed_count: 0,
total_tokens: Some(total_tokens),
cost_coverage: Default::default(),
costs: Vec::new(),
last_used_at: Some(Utc::now()),
request_buckets: Vec::new(),
models: Vec::new(),
}
}
pub(super) fn account_record(kind: &str) -> AccountRecord {
let now = Utc::now();
AccountRecord {
backend/crates/gateway-admin/tests/use_case/mod.rs
+54 -4
View File
@@ -7,7 +7,10 @@ mod settings;
mod system;
mod xai;
use std::sync::{Arc, Mutex};
use std::{
str::FromStr,
sync::{Arc, Mutex},
};
use async_trait::async_trait;
use futures::future::BoxFuture;
@@ -16,8 +19,8 @@ use gateway_admin::{
model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -244,6 +247,13 @@ impl AccountStore for UnavailableStore {
Err(unavailable("account usage"))
}
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
Err(unavailable("account quota window usage"))
}
async fn list_credentials(
&self,
_: &ProviderKind,
@@ -389,6 +399,15 @@ impl ObservabilityStore for UnavailableStore {
Err(unavailable("usage trend"))
}
async fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::UsageCalculatedBillingFact>>
{
Err(unavailable("usage billing facts"))
}
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
Err(unavailable("usage records"))
}
@@ -454,6 +473,7 @@ impl SettingsStore for UnavailableStore {
struct UnavailableProvider {
kind: ProviderKind,
dashboard_profile: Option<DashboardWireProfile>,
calculated_billing: Option<gateway_admin::model::observability::CalculatedBillingBreakdown>,
}
impl UnavailableProvider {
@@ -461,6 +481,7 @@ impl UnavailableProvider {
Self {
kind: ProviderKind::new(kind).expect("provider kind"),
dashboard_profile: None,
calculated_billing: None,
}
}
}
@@ -492,7 +513,7 @@ impl ProviderAdmin for UnavailableProvider {
Option<gateway_admin::model::observability::CalculatedBillingBreakdown>,
ProviderAdminError,
> {
Ok(None)
Ok(self.calculated_billing.clone())
}
async fn prepare_import(
@@ -587,6 +608,35 @@ pub(super) fn dashboard_profile_provider() -> Arc<dyn ProviderAdmin> {
error: None,
}),
}),
calculated_billing: None,
})
}
pub(super) fn calculated_billing_provider() -> Arc<dyn ProviderAdmin> {
let amount = |value| gateway_admin::model::observability::CurrencyCost {
currency: "USD".to_owned(),
amount: gateway_admin::model::observability::DecimalAmount::from_str(value)
.expect("test billing amount"),
};
Arc::new(UnavailableProvider {
kind: ProviderKind::new("openai").expect("provider kind"),
dashboard_profile: None,
calculated_billing: Some(
gateway_admin::model::observability::CalculatedBillingBreakdown {
input_amount: amount("0.8"),
output_amount: amount("0.2"),
cache_read_amount: amount("0"),
cache_write_amount: amount("0"),
standard_amount: amount("1"),
total_amount: amount("1.25"),
input_price_per_million: amount("0"),
output_price_per_million: amount("0"),
cache_read_price_per_million: amount("0"),
cache_write_price_per_million: amount("0"),
service_tier: None,
multiplier_percent: 125,
},
),
})
}
backend/crates/gateway-admin/tests/use_case/observability.rs
+64 -2
View File
@@ -14,7 +14,8 @@ use gateway_admin::{
AccountPoolMetrics, AttemptMetrics, CostCoverage, CurrencyCost, DashboardObservation,
DashboardRuntimeSlots, DiagnosticDimension, DiagnosticObservation, Granularity,
HealthStatus, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, RequestMetrics,
TimeRange, TrendKind, UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
TimeRange, TrendKind, UsageCalculatedBillingFact, UsageDetail, UsageFilter,
UsageOverview, UsagePage, UsageQuery,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -242,8 +243,22 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
.expect("bucket USD cost"),
}],
}]);
store.replace_calculated_billing_facts(vec![UsageCalculatedBillingFact {
bucket_start: quarter_hour_start(now),
provider_kind: "openai".to_owned(),
upstream_model_id: "gpt-5.5".to_owned(),
input_tokens: Some(800),
output_tokens: Some(200),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: gateway_admin::model::observability::DecimalAmount::from_str("1.25")
.expect("calculated total"),
},
}]);
store.replace_diagnostics(vec![diagnostic("openai", 3), diagnostic("xai", 1)]);
let services = observability_services(store).await;
let services = observability_services_with_calculated_billing(store).await;
let insights = services
.observability()
@@ -284,6 +299,21 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
.map(gateway_admin::model::observability::DecimalAmount::as_str),
Some("0.25")
);
assert_eq!(
insights
.cost
.standard_cost
.as_ref()
.map(gateway_admin::model::observability::DecimalAmount::as_str),
Some("1")
);
assert_eq!(
insights.cost.points[0]
.standard_cost
.as_ref()
.map(gateway_admin::model::observability::DecimalAmount::as_str),
Some("1")
);
assert_eq!(diagnostics.items[0].request_share, 0.75);
assert_eq!(diagnostics.items[1].request_share, 0.25);
}
@@ -291,6 +321,7 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
struct FixtureObservabilityStore {
trend: Mutex<Vec<RequestMetricPoint>>,
overview: Mutex<UsageOverview>,
calculated_billing_facts: Mutex<Vec<UsageCalculatedBillingFact>>,
diagnostics: Mutex<Vec<DiagnosticObservation>>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
}
@@ -305,6 +336,7 @@ impl FixtureObservabilityStore {
attempts: AttemptMetrics::default(),
providers: Vec::new(),
}),
calculated_billing_facts: Mutex::new(Vec::new()),
diagnostics: Mutex::new(Vec::new()),
runtime_slots: Mutex::new(None),
}
@@ -318,6 +350,13 @@ impl FixtureObservabilityStore {
*self.overview.lock().expect("overview") = overview;
}
fn replace_calculated_billing_facts(&self, facts: Vec<UsageCalculatedBillingFact>) {
*self
.calculated_billing_facts
.lock()
.expect("calculated billing facts") = facts;
}
fn replace_diagnostics(&self, diagnostics: Vec<DiagnosticObservation>) {
*self.diagnostics.lock().expect("diagnostics") = diagnostics;
}
@@ -360,6 +399,18 @@ impl ObservabilityStore for FixtureObservabilityStore {
Ok(self.trend.lock().expect("trend").clone())
}
async fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<UsageCalculatedBillingFact>> {
Ok(self
.calculated_billing_facts
.lock()
.expect("calculated billing facts")
.clone())
}
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
Err(super::unavailable("usage records"))
}
@@ -445,6 +496,17 @@ async fn observability_services(store: Arc<FixtureObservabilityStore>) -> AdminS
.await
}
async fn observability_services_with_calculated_billing(
store: Arc<FixtureObservabilityStore>,
) -> AdminServices {
super::AdminHarness::new()
.observability(store)
.settings(Arc::new(FixtureSettingsStore))
.provider(super::calculated_billing_provider())
.build()
.await
}
fn observation_range(end: DateTime<Utc>) -> TimeRange {
TimeRange::new(end - Duration::hours(24), end).expect("observation range")
}
backend/crates/gateway-api/src/admin/observability.rs
+52 -3
View File
@@ -1,6 +1,6 @@
//! Dashboard、用量与诊断查询的 wire 映射和固定路由。
use std::num::NonZeroU32;
use std::{collections::BTreeMap, num::NonZeroU32};
use axum::{
Router,
@@ -13,6 +13,7 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Duration, Timelike as _, Utc};
use gateway_admin::model::{PageSize as DomainPageSize, observability as domain};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError,
@@ -658,10 +659,20 @@ pub struct UsageRecordMetadataView {
pub response_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_request_id: Option<String>,
pub websocket_pool: Option<String>,
pub websocket_pool: Option<WebSocketPoolMetadataView>,
pub image_generation_requested: bool,
pub image_generation_succeeded: Option<bool>,
pub latency_details: UsageLatencyDetailsView,
/// Provider 安全观测直接展开，避免 API 层理解任何 Provider 专属字段。
#[serde(flatten)]
pub provider_metadata: BTreeMap<String, Value>,
}
/// WebSocket 池决策的稳定形状；与 Provider 选择逻辑无关。
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct WebSocketPoolMetadataView {
pub kind: String,
}
/// 逻辑请求在上游和输出阶段测得的时延事实。
@@ -1759,7 +1770,10 @@ fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView {
upstream_status_code: record.upstream_status_code.map(i64::from),
response_id: record.client_response_id.clone(),
upstream_request_id: record.upstream_request_id.clone(),
websocket_pool: record.websocket_pool.clone(),
websocket_pool: record
.websocket_pool
.clone()
.map(|kind| WebSocketPoolMetadataView { kind }),
image_generation_requested: record.image_generation_requested,
image_generation_succeeded: record.image_generation_succeeded,
latency_details: UsageLatencyDetailsView {
@@ -1772,6 +1786,7 @@ fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView {
first_token_ms: record.first_token_ms,
openai_processing_ms: record.provider_processing_ms,
},
provider_metadata: provider_metadata_fields(record.provider_metadata_json.as_deref()),
};
UsageRecordView {
id: record.id.clone(),
@@ -1824,6 +1839,40 @@ fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView {
}
}
fn provider_metadata_fields(value: Option<&str>) -> BTreeMap<String, Value> {
const CORE_METADATA_FIELDS: &[&str] = &[
"protocol",
"logicalOutcome",
"attemptCount",
"requestedModel",
"upstreamModel",
"clientIp",
"userAgent",
"reasoningEffort",
"reasoningPreset",
"compact",
"requestKind",
"subagentKind",
"transport",
"httpVersion",
"clientStatusCode",
"upstreamStatusCode",
"responseId",
"upstreamRequestId",
"websocketPool",
"imageGenerationRequested",
"imageGenerationSucceeded",
"latencyDetails",
];
let Some(Ok(Value::Object(fields))) = value.map(serde_json::from_str::<Value>) else {
return BTreeMap::new();
};
fields
.into_iter()
.filter(|(field, _)| !CORE_METADATA_FIELDS.contains(&field.as_str()))
.collect()
}
fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptView {
let occurred_at = attempt.occurred_at;
UsageAttemptView {
backend/crates/gateway-api/tests/admin/mod.rs
+18 -2
View File
@@ -14,8 +14,8 @@ use gateway_admin::{
model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -441,6 +441,13 @@ impl AccountStore for UnusedStore {
Err(unavailable("account usage"))
}
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
Err(unavailable("account quota window usage"))
}
async fn list_credentials(
&self,
_: &ProviderKind,
@@ -540,6 +547,15 @@ impl ObservabilityStore for UnusedStore {
Err(unavailable("usage trend"))
}
async fn usage_calculated_billing_facts(
&self,
_: TimeRange,
_: UsageFilter,
) -> AdminStoreResult<Vec<gateway_admin::model::observability::UsageCalculatedBillingFact>>
{
Err(unavailable("usage billing facts"))
}
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
let items = self.usage_records.lock().expect("usage records").clone();
Ok(UsagePage {
backend/crates/gateway-api/tests/admin/observability.rs
+15 -1
View File
@@ -326,6 +326,14 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
upstream_transport: Some("http_sse".to_owned()),
http_version: Some("h2".to_owned()),
websocket_pool: Some("reuse".to_owned()),
provider_metadata_json: Some(
serde_json::json!({
"effectiveModel": "grok-4.5",
"requestSummary": {"inputItemsCount": 2},
"transport": "must-not-overwrite-core",
})
.to_string(),
),
attempt_count: 1,
upstream_send_state: "sent".to_owned(),
downstream_committed_at: None,
@@ -409,18 +417,21 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
"requestKind": value["data"]["items"][0]["metadata"]["requestKind"],
"subagentKind": value["data"]["items"][0]["metadata"]["subagentKind"],
"compact": value["data"]["items"][0]["metadata"]["compact"],
"transport": value["data"]["items"][0]["metadata"]["transport"],
"transportDecisionWaitMs": value["data"]["items"][0]["metadata"]["latencyDetails"]["transportDecisionWaitMs"],
"wsConnectMs": value["data"]["items"][0]["metadata"]["latencyDetails"]["wsConnectMs"],
"firstReasoningMs": value["data"]["items"][0]["metadata"]["latencyDetails"]["firstReasoningMs"],
"firstTextMs": value["data"]["items"][0]["metadata"]["latencyDetails"]["firstTextMs"],
"firstTokenMs": value["data"]["items"][0]["metadata"]["latencyDetails"]["firstTokenMs"],
"effectiveModel": value["data"]["items"][0]["metadata"]["effectiveModel"],
"requestSummary": value["data"]["items"][0]["metadata"]["requestSummary"],
}),
serde_json::json!({
"route": "/v1/responses",
"authenticationKind": "oauth",
"imageInputTokens": 31,
"imageOutputTokens": 9,
"websocketPool": "reuse",
"websocketPool": {"kind": "reuse"},
"imageGenerationRequested": true,
"imageGenerationSucceeded": true,
"requestedModel": "grok-4.5",
@@ -432,11 +443,14 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
"requestKind": "review",
"subagentKind": "worker",
"compact": true,
"transport": "http_sse",
"transportDecisionWaitMs": 7,
"wsConnectMs": 11,
"firstReasoningMs": 19,
"firstTextMs": 23,
"firstTokenMs": 19,
"effectiveModel": "grok-4.5",
"requestSummary": {"inputItemsCount": 2},
})
);
}
backend/crates/gateway-core/src/engine/coordinator.rs
+40 -11
View File
@@ -3,7 +3,7 @@
use std::collections::{BTreeSet, VecDeque};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::time::{Duration, Instant, SystemTime};
use crate::accounting::{CostEstimate, CostSource, Usage};
use crate::engine::continuation::{
@@ -58,7 +58,7 @@ where
) -> Result<ResponseExecutionSession<S>, EngineError> {
let request_id = request.id.clone();
let client_api_key_ref = request.client_api_key_ref.clone();
let request_started_at = request.started_at;
let timing_started_at = Instant::now();
let deadline = request.deadline_at;
let account_state_owner = continuation
.as_ref()
@@ -77,7 +77,7 @@ where
engine: Arc::clone(&self.engine),
request_id,
client_api_key_ref,
request_started_at,
timing_started_at,
deadline,
pending_request: Some(request),
request_persisted: false,
@@ -155,7 +155,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
engine: Arc<GatewayEngine<S>>,
request_id: ModelRequestId,
client_api_key_ref: crate::policy::ClientApiKeyId,
request_started_at: SystemTime,
timing_started_at: Instant,
deadline: SystemTime,
pending_request: Option<NewModelRequest>,
request_persisted: bool,
@@ -499,7 +499,8 @@ where
.and_then(NonZeroU32::new)
.ok_or(EngineError::EmptyRoutingPlan)?;
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone()),
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
.with_timing_started_at(self.timing_started_at),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
@@ -692,8 +693,7 @@ where
}
async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> {
let observed_at = SystemTime::now();
let elapsed = elapsed_ms(self.request_started_at, observed_at);
let elapsed = elapsed_ms(self.timing_started_at);
observe_event_timing(&mut self.timings, event, elapsed);
if let GatewayEvent::Usage(observed) = event {
self.usage.merge(observed);
@@ -758,6 +758,18 @@ where
if let Some(value) = observed.first_event_ms {
self.timings.first_event_ms = Some(value);
}
if let Some(value) = observed.first_reasoning_ms {
self.timings.first_reasoning_ms = Some(value);
}
if let Some(value) = observed.first_text_ms {
self.timings.first_text_ms = Some(value);
}
if let Some(value) = observed.first_token_ms {
self.timings.first_token_ms = Some(value);
}
if let Some(value) = observed.provider_processing_ms {
self.timings.provider_processing_ms = Some(value);
}
current.response_observation = Some(observation);
Ok(())
}
@@ -949,7 +961,7 @@ where
return Ok(());
}
let completed_at = SystemTime::now();
self.timings.latency_ms = Some(elapsed_ms(self.request_started_at, completed_at));
self.timings.latency_ms = Some(elapsed_ms(self.timing_started_at));
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
.response_observation
@@ -974,6 +986,7 @@ where
.or(Some(200));
let (upstream_transport, http_version, websocket_pool) =
self.current_transport_observation();
let provider_metadata_json = self.current_provider_metadata_json();
self.engine
.store()
.finalize_model_request(ModelRequestFinalization {
@@ -990,6 +1003,7 @@ where
upstream_transport,
http_version,
websocket_pool,
provider_metadata_json,
error: None,
provider_error_code: None,
retry_after_ms: None,
@@ -1096,7 +1110,7 @@ where
return Ok(());
}
let completed_at = SystemTime::now();
self.timings.latency_ms = Some(elapsed_ms(self.request_started_at, completed_at));
self.timings.latency_ms = Some(elapsed_ms(self.timing_started_at));
let upstream_request_id = self.current.as_ref().and_then(|current| {
current
.response_observation
@@ -1120,6 +1134,7 @@ where
.and_then(ProviderResponseObservation::status_code);
let (upstream_transport, http_version, websocket_pool) =
self.current_transport_observation();
let provider_metadata_json = self.current_provider_metadata_json();
self.engine
.store()
.finalize_model_request(ModelRequestFinalization {
@@ -1138,6 +1153,7 @@ where
upstream_transport,
http_version,
websocket_pool,
provider_metadata_json,
error: Some(finalization.error),
provider_error_code: finalization.provider_error_code,
retry_after_ms: finalization.retry_after_ms,
@@ -1170,6 +1186,14 @@ where
)
}
fn current_provider_metadata_json(&self) -> Option<String> {
self.current
.as_ref()
.and_then(|current| current.response_observation.as_ref())
.and_then(ProviderResponseObservation::provider_metadata)
.map(|metadata| metadata.as_json().to_owned())
}
fn image_generation_succeeded(&self) -> Option<bool> {
self.image_generation_requested
.then(|| self.usage.image_output_tokens.unwrap_or_default() > 0)
@@ -1325,13 +1349,18 @@ fn observe_event_timing(timings: &mut ModelRequestTimings, event: &GatewayEvent,
timings.first_text_ms.get_or_insert(elapsed_ms);
timings.first_token_ms.get_or_insert(elapsed_ms);
}
// `response.output_item.added` 会先投影一个空参数的 tool delta；它只是结构帧，
// 不能抢在真实工具参数之前成为首个可消费 token。
GatewayEvent::ToolCallDelta(delta) if !delta.arguments_delta.is_empty() => {
timings.first_token_ms.get_or_insert(elapsed_ms);
}
GatewayEvent::CalculatedCost(_) | GatewayEvent::ProviderCost(_) => {}
_ => {}
}
}
fn elapsed_ms(started_at: SystemTime, observed_at: SystemTime) -> u64 {
duration_ms(observed_at.duration_since(started_at).unwrap_or_default())
fn elapsed_ms(started_at: Instant) -> u64 {
duration_ms(started_at.elapsed())
}
fn duration_ms(duration: Duration) -> u64 {
backend/crates/gateway-core/src/engine/mod.rs
+25 -2
View File
@@ -16,7 +16,7 @@ use std::net::IpAddr;
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use futures::channel::oneshot;
@@ -332,17 +332,26 @@ pub enum ContinuationAttempt {
pub struct RequestAttemptContext {
request_id: ModelRequestId,
client_api_key_ref: ClientApiKeyId,
timing_started_at: Instant,
}
impl RequestAttemptContext {
#[must_use]
pub const fn new(request_id: ModelRequestId, client_api_key_ref: ClientApiKeyId) -> Self {
pub fn new(request_id: ModelRequestId, client_api_key_ref: ClientApiKeyId) -> Self {
Self {
request_id,
client_api_key_ref,
timing_started_at: Instant::now(),
}
}
/// 覆盖本次请求的单调计时原点。
#[must_use]
pub fn with_timing_started_at(mut self, timing_started_at: Instant) -> Self {
self.timing_started_at = timing_started_at;
self
}
#[must_use]
pub const fn request_id(&self) -> &ModelRequestId {
&self.request_id
@@ -352,6 +361,12 @@ impl RequestAttemptContext {
pub const fn client_api_key_ref(&self) -> &ClientApiKeyId {
&self.client_api_key_ref
}
/// 返回贯穿本次请求的单调计时原点。
#[must_use]
pub const fn timing_started_at(&self) -> Instant {
self.timing_started_at
}
}
/// Provider 每次执行可见的 request-local context。
@@ -416,6 +431,12 @@ impl AttemptContext {
self.request.client_api_key_ref()
}
/// 返回本次请求统一的单调计时原点。
#[must_use]
pub const fn timing_started_at(&self) -> Instant {
self.request.timing_started_at()
}
#[must_use]
pub const fn attempt_index(&self) -> NonZeroU32 {
self.attempt_index
@@ -552,6 +573,8 @@ pub struct ModelRequestFinalization {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
/// Provider 已筛选的专有观测 JSON；Core 不解释字段。
pub provider_metadata_json: Option<String>,
pub error: Option<GatewayError>,
pub provider_error_code: Option<String>,
pub retry_after_ms: Option<u64>,
backend/crates/gateway-core/src/engine/provider.rs
+3 -1
View File
@@ -273,7 +273,9 @@ impl ProviderStreamAccountFeedback {
fn report_failure(&mut self, error: &ProviderError) {
if self.reported
|| error.send_state() == UpstreamSendState::NotSent
// `Ambiguous` 仍须关闭重放边界，但不能证明失败由账号造成；将它计入
// 账号评分会把传输不确定性错误归因给账号。
|| error.send_state() != UpstreamSendState::Sent
|| matches!(
error.kind(),
ProviderErrorKind::Cancelled | ProviderErrorKind::ProcessTerminated
backend/crates/gateway-core/src/event.rs
+55
View File
@@ -398,6 +398,46 @@ pub struct ProviderResponseTimings {
pub connect_ms: Option<u64>,
pub headers_ms: Option<u64>,
pub first_event_ms: Option<u64>,
pub first_reasoning_ms: Option<u64>,
pub first_text_ms: Option<u64>,
pub first_token_ms: Option<u64>,
pub provider_processing_ms: Option<u64>,
}
/// Provider 已筛选的响应观测 JSON。
///
/// Core 只验证它是有界 JSON object，并在请求终态原样交给存储层；字段语义、
/// 脱敏和版本演进均由所属 Provider 负责。这样 Provider 专有协议不会渗入路由
/// 或管理领域。
#[derive(Clone, PartialEq, Eq)]
pub struct ProviderResponseMetadata(String);
impl ProviderResponseMetadata {
const MAX_BYTES: usize = 32 * 1024;
/// 从 Provider 已筛选的 JSON object 创建观测快照。
#[must_use]
pub fn new(json: String) -> Option<Self> {
if json.len() > Self::MAX_BYTES {
return None;
}
serde_json::from_str::<Value>(&json)
.ok()
.filter(Value::is_object)
.map(|_| Self(json))
}
/// 返回不透明的 JSON 文本，供持久化 adapter 原样验证和写入。
#[must_use]
pub fn as_json(&self) -> &str {
&self.0
}
}
impl fmt::Debug for ProviderResponseMetadata {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("ProviderResponseMetadata(<provider-owned>)")
}
}
/// Provider 已筛选、可由协议 adapter 再次按白名单表达的响应头。
@@ -451,6 +491,7 @@ pub struct ProviderResponseObservation {
request_id: Option<SafeUpstreamValue>,
timings: ProviderResponseTimings,
client_headers: Vec<ProviderResponseHeader>,
provider_metadata: Option<ProviderResponseMetadata>,
}
impl ProviderResponseObservation {
@@ -464,6 +505,7 @@ impl ProviderResponseObservation {
request_id: None,
timings: ProviderResponseTimings::default(),
client_headers: Vec::new(),
provider_metadata: None,
}
}
@@ -503,6 +545,13 @@ impl ProviderResponseObservation {
self
}
/// 附加由 Provider 自己定义且已筛选的响应观测快照。
#[must_use]
pub fn with_provider_metadata(mut self, metadata: ProviderResponseMetadata) -> Self {
self.provider_metadata = Some(metadata);
self
}
#[must_use]
pub const fn transport(&self) -> &UpstreamTransport {
&self.transport
@@ -537,6 +586,12 @@ impl ProviderResponseObservation {
pub fn client_headers(&self) -> &[ProviderResponseHeader] {
&self.client_headers
}
/// 返回 Provider 专有的安全观测 JSON；Core 不读取其字段。
#[must_use]
pub const fn provider_metadata(&self) -> Option<&ProviderResponseMetadata> {
self.provider_metadata.as_ref()
}
}
/// WebSocket 请求使用新连接或复用池中连接的事实。
backend/crates/gateway-core/tests/engine/coordinator.rs
+124 -2
View File
@@ -29,8 +29,9 @@ use gateway_core::error::{
ContinuationFailure, ProviderError, ProviderErrorKind, SafeUpstreamValue, StoreError,
};
use gateway_core::event::{
GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseObservation,
ProviderResponseTimings, ResponseMeta, UpstreamHttpVersion, WebSocketPoolKind,
ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent,
ProviderResponseObservation, ProviderResponseTimings, ResponseMeta, ToolCallDelta,
UpstreamHttpVersion, WebSocketPoolKind,
};
use gateway_core::operation::{
EmbedRequest, GenerateRequest, Operation, ProtocolPayload, ProviderSessionState, RetrySafety,
@@ -67,6 +68,11 @@ struct FinalState {
connect_ms: Option<u64>,
headers_ms: Option<u64>,
first_event_ms: Option<u64>,
first_reasoning_ms: Option<u64>,
first_text_ms: Option<u64>,
first_token_ms: Option<u64>,
provider_processing_ms: Option<u64>,
provider_metadata_json: Option<String>,
cost_source: CostSource,
cost_ticks: Option<u128>,
}
@@ -191,6 +197,11 @@ impl ExecutionStore for FakeStore {
connect_ms: finalization.timings.connect_ms,
headers_ms: finalization.timings.headers_ms,
first_event_ms: finalization.timings.first_event_ms,
first_reasoning_ms: finalization.timings.first_reasoning_ms,
first_text_ms: finalization.timings.first_text_ms,
first_token_ms: finalization.timings.first_token_ms,
provider_processing_ms: finalization.timings.provider_processing_ms,
provider_metadata_json: finalization.provider_metadata_json,
cost_source: finalization.cost.source(),
cost_ticks: finalization
.cost
@@ -632,6 +643,7 @@ fn response_observation_is_persisted_but_never_delivered() {
connect_ms: Some(11),
headers_ms: Some(13),
first_event_ms: Some(17),
..ProviderResponseTimings::default()
});
let mut items = vec![Ok(ProviderEvent::observation(observation))];
items.extend(
@@ -673,6 +685,115 @@ fn response_observation_is_persisted_but_never_delivered() {
assert_eq!(finalization.first_event_ms, Some(17));
}
#[test]
fn response_observation_should_persist_provider_metadata_and_output_timings() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation);
let metadata = gateway_core::event::ProviderResponseMetadata::new(
serde_json::json!({"effectiveModel": "gpt-test", "firstTokenMs": 23}).to_string(),
)
.expect("provider metadata object");
let observation = ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("actual transport"),
)
.with_timings(ProviderResponseTimings {
first_event_ms: Some(11),
first_reasoning_ms: Some(17),
first_text_ms: Some(19),
first_token_ms: Some(17),
provider_processing_ms: Some(5),
..ProviderResponseTimings::default()
})
.with_provider_metadata(metadata);
let mut items = vec![Ok(ProviderEvent::observation(observation))];
items.extend(
complete_stream(Some(21))
.into_iter()
.map(|event| event.map(ProviderEvent::canonical)),
);
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
account_id: "acct_observed",
items,
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("collect response");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let state = store.state.lock().expect("store lock");
let finalization = &state.finalizations[0];
assert_eq!(finalization.first_event_ms, Some(11));
assert_eq!(finalization.first_reasoning_ms, Some(17));
assert_eq!(finalization.first_text_ms, Some(19));
assert_eq!(finalization.first_token_ms, Some(17));
assert_eq!(finalization.provider_processing_ms, Some(5));
assert_eq!(
finalization.provider_metadata_json.as_deref(),
Some("{\"effectiveModel\":\"gpt-test\",\"firstTokenMs\":23}")
);
}
#[test]
fn empty_tool_call_delta_should_not_preempt_provider_first_token_timing() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation);
let transport = UpstreamTransport::new("websocket").expect("actual transport");
let initial_observation = ProviderResponseObservation::new(transport.clone());
let output_observation =
ProviderResponseObservation::new(transport).with_timings(ProviderResponseTimings {
first_token_ms: Some(37),
..ProviderResponseTimings::default()
});
let items = vec![
Ok(ProviderEvent::observation(initial_observation)),
Ok(ProviderEvent::canonical(GatewayEvent::Started(
ResponseMeta::new("response-tool", "gpt-5"),
))),
Ok(ProviderEvent::canonical(GatewayEvent::ContentAdded(
ContentItem::new(0, ContentKind::ToolCall),
))),
Ok(ProviderEvent::canonical(GatewayEvent::ToolCallDelta(
ToolCallDelta {
content_index: 0,
call_id: "call-tool".to_owned(),
name: Some("apply_patch".to_owned()),
arguments_delta: String::new(),
},
))),
Ok(ProviderEvent::observation(output_observation)),
Ok(ProviderEvent::canonical(GatewayEvent::Completed(
ResponseMeta::new("response-tool", "gpt-5"),
))),
];
let (coordinator, store, _) = coordinator(vec![Script::ObservedStream {
account_id: "acct_tool",
items,
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("collect response");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let state = store.state.lock().expect("store lock");
assert_eq!(state.finalizations[0].first_token_ms, Some(37));
}
#[test]
fn unknown_wire_event_before_response_identity_is_discarded_with_retried_attempt() {
let operation = operation(RetrySafety::Idempotent);
@@ -734,6 +855,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
connect_ms: Some(11),
headers_ms: Some(13),
first_event_ms: Some(987_654),
..ProviderResponseTimings::default()
});
let second_observation = ProviderResponseObservation::new(
UpstreamTransport::new("http_sse").expect("second transport"),
backend/crates/gateway-core/tests/engine/provider.rs
+7 -1
View File
@@ -203,13 +203,15 @@ fn provider_stream_should_not_sequence_gate_canonical_observation_attached_to_wi
}
#[test]
fn provider_stream_should_report_sent_failure_but_ignore_not_sent_failure() {
fn provider_stream_should_report_confirmed_sent_failure_but_ignore_unconfirmed_failure() {
let feedback = Arc::new(AccountFeedbackStats::default());
let provider = ProviderKind::new("xai").expect("valid provider");
let sent_account = ProviderAccountId::new("acct_stream_sent").expect("account");
let ambiguous_account = ProviderAccountId::new("acct_stream_ambiguous").expect("account");
let not_sent_account = ProviderAccountId::new("acct_stream_not_sent").expect("account");
for (account, send_state) in [
(sent_account.clone(), UpstreamSendState::Sent),
(ambiguous_account.clone(), UpstreamSendState::Ambiguous),
(not_sent_account.clone(), UpstreamSendState::NotSent),
] {
let metadata = ProviderCallMetadata::new(
@@ -241,6 +243,10 @@ fn provider_stream_should_report_sent_failure_but_ignore_not_sent_failure() {
feedback.scheduling_signals(&provider, &sent_account).0,
Some(2_000)
);
assert_eq!(
feedback.scheduling_signals(&provider, &ambiguous_account),
(None, None)
);
assert_eq!(
feedback.scheduling_signals(&provider, &not_sent_account),
(None, None)
backend/crates/gateway-host/tests/system_update/mod.rs
+9 -1
View File
@@ -49,7 +49,8 @@ async fn restart_should_request_process_restart_inside_docker() {
#[tokio::test]
async fn restart_should_spawn_replacement_before_shutdown_outside_docker() {
let fixture = Fixture::new();
fixture.write_executable("#!/bin/sh\nexit 0\n");
let replacement_marker = fixture.root.path().join("replacement-ran");
fixture.write_executable("#!/bin/sh\n: > \"${0%/*}/replacement-ran\"\n");
let mut config = fixture.config("http://127.0.0.1:1/repos");
config.self_restart_enabled = true;
let shutdown = CancellationToken::new();
@@ -66,6 +67,13 @@ async fn restart_should_spawn_replacement_before_shutdown_outside_docker() {
tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled())
.await
.expect("shutdown requested after spawn");
tokio::time::timeout(Duration::from_secs(2), async {
while !replacement_marker.is_file() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("replacement executed before fixture cleanup");
}
#[tokio::test]
backend/crates/gateway-store/src/postgres/execution.rs
+25 -1
View File
@@ -2,6 +2,7 @@
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde_json::Value;
use sqlx::PgPool;
use uuid::Uuid;
@@ -223,6 +224,7 @@ pub struct ModelRequestFinalization {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub provider_metadata_json: Option<Value>,
pub error_kind: Option<String>,
pub provider_error_code: Option<String>,
pub error_message: Option<String>,
@@ -274,6 +276,13 @@ impl ModelRequestFinalization {
{
return Err(invalid("websocket pool must be new or reuse"));
}
if self
.provider_metadata_json
.as_ref()
.is_some_and(|metadata| !metadata.is_object())
{
return Err(invalid("provider observation must be a JSON object"));
}
self.timings.validate()
}
}
@@ -498,7 +507,8 @@ impl ModelRequestRepository for PgExecutionStore {
first_text_ms = $32, first_token_ms = $33, provider_processing_ms = $34,
latency_ms = $35, completed_at = $36,
upstream_transport = coalesce($37, upstream_transport),
http_version = coalesce($38, http_version), websocket_pool = $39
http_version = coalesce($38, http_version), websocket_pool = $39,
provider_observation_json = $40
where id = $1 and outcome = 'running'",
)
.bind(&finalization.model_request_id)
@@ -585,6 +595,7 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(finalization.upstream_transport)
.bind(finalization.http_version)
.bind(finalization.websocket_pool)
.bind(finalization.provider_metadata_json.map(sqlx::types::Json))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("finalize model request"))?;
@@ -769,6 +780,18 @@ impl ExecutionStore for PgExecutionStore {
&self,
finalization: CoreModelRequestFinalization,
) -> Result<(), CoreStoreError> {
let provider_metadata_json = finalization
.provider_metadata_json
.as_deref()
.map(serde_json::from_str::<Value>)
.transpose()
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
if provider_metadata_json
.as_ref()
.is_some_and(|metadata| !metadata.is_object())
{
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
}
let (cost_source, cost_amount, cost_currency) = match finalization.cost.total() {
Some(total) => (
cost_source_from_core(finalization.cost.source()),
@@ -809,6 +832,7 @@ impl ExecutionStore for PgExecutionStore {
upstream_transport: finalization.upstream_transport,
http_version: finalization.http_version,
websocket_pool: finalization.websocket_pool,
provider_metadata_json,
error_kind,
provider_error_code: finalization.provider_error_code,
error_message,
backend/crates/gateway-store/src/postgres/observability.rs
+203 -48
View File
@@ -364,6 +364,19 @@ pub struct RequestMetricPoint {
pub costs: Vec<CurrencyCostTotal>,
}
/// 已完整交付且由 Provider 计算费用的请求事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CalculatedUsageBillingFact {
pub bucket_start: DateTime<Utc>,
pub provider_kind: String,
pub upstream_model_id: String,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
pub cache_write_tokens: Option<u64>,
pub total: CurrencyCostTotal,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderAccountMetrics {
pub total: u64,
@@ -509,6 +522,7 @@ pub struct UsageRecord {
pub upstream_transport: Option<String>,
pub http_version: Option<String>,
pub websocket_pool: Option<String>,
pub provider_metadata_json: Option<String>,
pub attempt_count: u32,
pub upstream_send_state: String,
pub downstream_committed_at: Option<DateTime<Utc>>,
@@ -679,6 +693,11 @@ pub trait ObservabilityRepository: Send + Sync {
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<RequestMetricPoint>>;
async fn usage_calculated_billing_facts(
&self,
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>>;
async fn provider_account_usage(
&self,
query: ProviderAccountUsageQuery,
@@ -793,6 +812,14 @@ impl ObservabilityRepository for PgObservabilityRepository {
request_metric_series(&self.pool, range, &filter).await
}
async fn usage_calculated_billing_facts(
&self,
range: ObservabilityRange,
filter: UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>> {
calculated_usage_billing_facts(&self.pool, range, &filter).await
}
async fn provider_account_usage(
&self,
query: ProviderAccountUsageQuery,
@@ -919,6 +946,20 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
.collect()
}
async fn usage_calculated_billing_facts(
&self,
range: admin_observability::TimeRange,
filter: admin_observability::UsageFilter,
) -> AdminStoreResult<Vec<admin_observability::UsageCalculatedBillingFact>> {
self.repository
.usage_calculated_billing_facts(store_range(range)?, store_usage_filter(filter))
.await
.map_err(observability_error)?
.into_iter()
.map(admin_calculated_usage_billing_fact)
.collect()
}
async fn list_usage_records(
&self,
query: admin_observability::UsageQuery,
@@ -1330,6 +1371,24 @@ fn admin_request_metric_point(
})
}
fn admin_calculated_usage_billing_fact(
fact: CalculatedUsageBillingFact,
) -> AdminStoreResult<admin_observability::UsageCalculatedBillingFact> {
Ok(admin_observability::UsageCalculatedBillingFact {
bucket_start: fact.bucket_start,
provider_kind: fact.provider_kind,
upstream_model_id: fact.upstream_model_id,
input_tokens: fact.input_tokens,
output_tokens: fact.output_tokens,
cached_tokens: fact.cached_tokens,
cache_write_tokens: fact.cache_write_tokens,
total: admin_observability::CurrencyCost {
currency: fact.total.currency,
amount: admin_decimal_amount(fact.total.amount)?,
},
})
}
const fn admin_granularity(
granularity: ObservationGranularity,
) -> admin_observability::Granularity {
@@ -1394,6 +1453,7 @@ fn admin_usage_record(record: UsageRecord) -> AdminStoreResult<admin_observabili
upstream_transport: record.upstream_transport,
http_version: record.http_version,
websocket_pool: record.websocket_pool,
provider_metadata_json: record.provider_metadata_json,
attempt_count: record.attempt_count,
upstream_send_state: record.upstream_send_state,
downstream_committed_at: record.downstream_committed_at,
@@ -1629,6 +1689,7 @@ async fn request_metrics(
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
let row = query
.build()
@@ -1695,6 +1756,7 @@ async fn request_metric_series(
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" group by bucket_start order by bucket_start");
let rows = query
@@ -1726,6 +1788,44 @@ async fn request_metric_series(
fill_metric_gaps(range, granularity, points)
}
async fn calculated_usage_billing_facts(
pool: &PgPool,
range: ObservabilityRange,
filter: &UsageRecordFilter,
) -> StoreResult<Vec<CalculatedUsageBillingFact>> {
filter.validate()?;
let granularity = granularity_for(range);
let mut query = QueryBuilder::<Postgres>::new("select date_bin(");
query.push_bind(granularity.sql_interval());
query.push(
"::interval, mr.started_at, timestamptz '1970-01-01 00:00:00+00') as bucket_start,
mr.provider_kind, mr.upstream_model_id,
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
mr.cost_currency, mr.cost_amount::text as amount
from model_requests mr where mr.started_at >= ",
);
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
query.push(
" and mr.cost_source = 'calculated'
and mr.cost_amount is not null and mr.cost_currency is not null
and nullif(mr.provider_kind, '') is not null
and nullif(mr.upstream_model_id, '') is not null",
);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" order by bucket_start, mr.id");
query
.build()
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load calculated usage billing facts"))?
.iter()
.map(calculated_usage_billing_fact_from_row)
.collect()
}
async fn request_costs_by_bucket(
pool: &PgPool,
range: ObservabilityRange,
@@ -1744,6 +1844,7 @@ async fn request_costs_by_bucket(
query.push(" and mr.started_at < ");
query.push_bind(range.end);
query.push(" and mr.cost_amount is not null and mr.cost_currency is not null");
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" group by bucket_start, mr.cost_currency order by bucket_start, mr.cost_currency");
let rows = query
@@ -1776,7 +1877,12 @@ async fn attempt_metrics(
query.push_bind(range.end);
push_usage_filter(&mut query, filter, "mr");
query.push(
"), failures as (
"), usage_facts as (
select * from selected_requests
where outcome = 'succeeded'
and downstream_committed_at is not null
and client_status_code between 200 and 399
), failures as (
select coalesce(oe.occurrence_count, 1)::bigint as occurrences,
oe.failure_kind, oe.status_code
from ops_events oe
@@ -1809,13 +1915,13 @@ async fn attempt_metrics(
coalesce((select sum(occurrences) from failures
where status_code between 500 and 599), 0)::bigint
as provider_5xx_count,
coalesce((select count(*) from selected_requests
coalesce((select count(*) from usage_facts
where cost_source = 'provider_reported'), 0)::bigint
as provider_reported_count,
coalesce((select count(*) from selected_requests
coalesce((select count(*) from usage_facts
where cost_source = 'calculated'), 0)::bigint
as calculated_count,
coalesce((select count(*) from selected_requests
coalesce((select count(*) from usage_facts
where cost_source = 'unavailable'), 0)::bigint
as unavailable_count",
);
@@ -1851,6 +1957,7 @@ async fn request_costs(
query.push(" and mr.started_at < ");
query.push_bind(range.end);
query.push(" and mr.cost_amount is not null and mr.cost_currency is not null");
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" group by mr.cost_currency order by mr.cost_currency");
query
@@ -1938,6 +2045,7 @@ async fn provider_observations(
query.push_bind(range.start);
query.push(" and mr.started_at < ");
query.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" group by coalesce(mr.provider_kind, 'unrouted') order by request_count desc");
let rows = query
@@ -2129,6 +2237,15 @@ fn push_usage_filter(query: &mut QueryBuilder<Postgres>, filter: &UsageRecordFil
}
}
/// 只把已完整交付给客户端的成功响应投影为使用事实。
///
/// `model_requests` 同时承担执行审计：包括上游已发送、但尚未收到首个事件就断开的
/// WebSocket 请求。这些失败仍必须留给 Ops Errors 和调度健康度分析，不能混入用量、
/// 成本、账号使用次数或请求明细。
fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres>, alias: &str) {
query.push(format!(" and {alias}.outcome = 'succeeded' and {alias}.downstream_committed_at is not null and {alias}.client_status_code between 200 and 399"));
}
async fn provider_account_usage(
pool: &PgPool,
query: ProviderAccountUsageQuery,
@@ -2171,6 +2288,7 @@ async fn provider_account_usage(
statement.push_bind(query.range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(query.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
if let Some(account_ids) = &query.account_ids {
statement.push(" where pa.id = any(");
statement.push_bind(account_ids.clone());
@@ -2263,11 +2381,14 @@ async fn provider_account_request_buckets(
"select provider_account_ref,
floor(extract(epoch from (started_at - $1)) / 3600)::bigint as bucket_index,
count(*)::bigint as request_count
from model_requests
where provider_account_ref = any($2::text[])
and started_at >= $1 and started_at < $3
group by provider_account_ref, bucket_index
order by provider_account_ref, bucket_index",
from model_requests mr
where mr.provider_account_ref = any($2::text[])
and mr.started_at >= $1 and mr.started_at < $3
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
group by mr.provider_account_ref, bucket_index
order by mr.provider_account_ref, bucket_index",
)
.bind(range.start)
.bind(account_ids)
@@ -2317,13 +2438,16 @@ async fn account_costs(
account_ids: &[String],
) -> StoreResult<HashMap<String, Vec<CurrencyCostTotal>>> {
let rows = sqlx::query(
"select provider_account_ref, cost_currency, sum(cost_amount)::text as amount
from model_requests
where provider_account_ref = any($1::text[])
and started_at >= $2 and started_at < $3
and cost_amount is not null and cost_currency is not null
group by provider_account_ref, cost_currency
order by provider_account_ref, cost_currency",
"select mr.provider_account_ref, mr.cost_currency, sum(mr.cost_amount)::text as amount
from model_requests mr
where mr.provider_account_ref = any($1::text[])
and mr.started_at >= $2 and mr.started_at < $3
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_amount is not null and mr.cost_currency is not null
group by mr.provider_account_ref, mr.cost_currency
order by mr.provider_account_ref, mr.cost_currency",
)
.bind(account_ids)
.bind(range.start)
@@ -2347,34 +2471,37 @@ async fn account_model_rows(
account_ids: &[String],
) -> StoreResult<Vec<sqlx::postgres::PgRow>> {
sqlx::query(
"select provider_account_ref,
coalesce(upstream_model_id, requested_model_id) as model,
"select mr.provider_account_ref,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
count(*)::bigint as request_count,
count(*) filter (where outcome = 'succeeded')::bigint as success_count,
sum(input_tokens)::bigint as input_tokens,
sum(output_tokens)::bigint as output_tokens,
sum(cached_tokens)::bigint as cached_tokens,
sum(cache_write_tokens)::bigint as cache_write_tokens,
sum(reasoning_tokens)::bigint as reasoning_tokens,
sum(image_input_tokens)::bigint as image_input_tokens,
sum(image_output_tokens)::bigint as image_output_tokens,
count(*) filter (where image_generation_succeeded is true)::bigint
count(*) filter (where mr.outcome = 'succeeded')::bigint as success_count,
sum(mr.input_tokens)::bigint as input_tokens,
sum(mr.output_tokens)::bigint as output_tokens,
sum(mr.cached_tokens)::bigint as cached_tokens,
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
sum(mr.image_input_tokens)::bigint as image_input_tokens,
sum(mr.image_output_tokens)::bigint as image_output_tokens,
count(*) filter (where mr.image_generation_succeeded is true)::bigint
as image_request_count,
count(*) filter (where image_generation_succeeded is false)::bigint
count(*) filter (where mr.image_generation_succeeded is false)::bigint
as image_request_failed_count,
sum(total_tokens)::bigint as total_tokens,
count(*) filter (where cost_source = 'provider_reported')::bigint
sum(mr.total_tokens)::bigint as total_tokens,
count(*) filter (where mr.cost_source = 'provider_reported')::bigint
as provider_reported_count,
count(*) filter (where cost_source = 'calculated')::bigint
count(*) filter (where mr.cost_source = 'calculated')::bigint
as calculated_count,
count(*) filter (where cost_source = 'unavailable')::bigint
count(*) filter (where mr.cost_source = 'unavailable')::bigint
as unavailable_count,
max(started_at) as last_used_at
from model_requests
where provider_account_ref = any($1::text[])
and started_at >= $2 and started_at < $3
group by provider_account_ref, coalesce(upstream_model_id, requested_model_id)
order by provider_account_ref, request_count desc, model",
max(mr.started_at) as last_used_at
from model_requests mr
where mr.provider_account_ref = any($1::text[])
and mr.started_at >= $2 and mr.started_at < $3
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id)
order by mr.provider_account_ref, request_count desc, model",
)
.bind(account_ids)
.bind(range.start)
@@ -2390,16 +2517,19 @@ async fn account_model_costs(
account_ids: &[String],
) -> StoreResult<HashMap<(String, String), Vec<CurrencyCostTotal>>> {
let rows = sqlx::query(
"select provider_account_ref,
coalesce(upstream_model_id, requested_model_id) as model,
cost_currency, sum(cost_amount)::text as amount
from model_requests
where provider_account_ref = any($1::text[])
and started_at >= $2 and started_at < $3
and cost_amount is not null and cost_currency is not null
group by provider_account_ref, coalesce(upstream_model_id, requested_model_id),
cost_currency
order by provider_account_ref, model, cost_currency",
"select mr.provider_account_ref,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
mr.cost_currency, sum(mr.cost_amount)::text as amount
from model_requests mr
where mr.provider_account_ref = any($1::text[])
and mr.started_at >= $2 and mr.started_at < $3
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_amount is not null and mr.cost_currency is not null
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id),
mr.cost_currency
order by mr.provider_account_ref, model, mr.cost_currency",
)
.bind(account_ids)
.bind(range.start)
@@ -2456,6 +2586,7 @@ const USAGE_RECORD_SELECT: &str =
pa.name as provider_account_name, pa.email as provider_account_email,
pa.authentication_kind as provider_account_authentication_kind,
mr.upstream_model_id, mr.upstream_transport, mr.http_version, mr.websocket_pool,
mr.provider_observation_json,
mr.attempt_count, mr.upstream_send_state, mr.downstream_committed_at,
mr.outcome, mr.client_status_code, mr.upstream_status_code,
mr.client_response_id, mr.upstream_request_id, mr.upstream_response_id,
@@ -2483,6 +2614,7 @@ async fn list_usage_records(
statement.push_bind(query.range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(query.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, &query.filter, "mr");
if let Some(cursor) = &query.cursor {
statement.push(" and (mr.started_at, mr.id) < (");
@@ -2540,6 +2672,7 @@ async fn count_usage_records(
statement.push_bind(range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, filter, "mr");
let total = statement
.build_query_scalar::<i64>()
@@ -2555,6 +2688,7 @@ async fn usage_record_detail(pool: &PgPool, request_id: &str) -> StoreResult<Usa
let mut statement = QueryBuilder::<Postgres>::new(USAGE_RECORD_SELECT);
statement.push(" where mr.id = ");
statement.push_bind(request_id.to_owned());
push_completed_usage_fact_filter(&mut statement, "mr");
let row = statement
.build()
.fetch_optional(pool)
@@ -2686,6 +2820,7 @@ async fn usage_diagnostics(
statement.push_bind(range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, filter, "mr");
statement.push(" group by dimension_name order by request_count desc, dimension_name limit ");
statement.push_bind(DIAGNOSTIC_LIMIT);
@@ -2740,6 +2875,7 @@ async fn diagnostic_costs(
statement.push(" and mr.started_at < ");
statement.push_bind(range.end);
statement.push(" and mr.cost_amount is not null and mr.cost_currency is not null");
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, filter, "mr");
statement.push(
" group by dimension_name, mr.cost_currency order by dimension_name, mr.cost_currency",
@@ -2943,6 +3079,10 @@ fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<UsageRecord
upstream_transport: get(row, "upstream_transport")?,
http_version: get(row, "http_version")?,
websocket_pool: get(row, "websocket_pool")?,
provider_metadata_json: get::<Option<serde_json::Value>>(row, "provider_observation_json")?
.map(|value| serde_json::to_string(&value))
.transpose()
.map_err(|_| postgres_unavailable("encode provider observation"))?,
attempt_count: to_u32(get(row, "attempt_count")?)?,
upstream_send_state: get(row, "upstream_send_state")?,
downstream_committed_at: get(row, "downstream_committed_at")?,
@@ -3026,6 +3166,21 @@ fn cost_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<CurrencyCostTotal>
})
}
fn calculated_usage_billing_fact_from_row(
row: &sqlx::postgres::PgRow,
) -> StoreResult<CalculatedUsageBillingFact> {
Ok(CalculatedUsageBillingFact {
bucket_start: get(row, "bucket_start")?,
provider_kind: get(row, "provider_kind")?,
upstream_model_id: get(row, "upstream_model_id")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
output_tokens: optional_unsigned(row, "output_tokens")?,
cached_tokens: optional_unsigned(row, "cached_tokens")?,
cache_write_tokens: optional_unsigned(row, "cache_write_tokens")?,
total: cost_from_row(row)?,
})
}
fn optional_decimal(
row: &sqlx::postgres::PgRow,
column: &'static str,
backend/crates/gateway-store/src/postgres/provider_accounts.rs
+191 -3
View File
@@ -16,8 +16,8 @@ use gateway_admin::{
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, DeleteAccounts, SetAccountEnabled,
SortDirection as AdminSortDirection,
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
DeleteAccounts, SetAccountEnabled, SortDirection as AdminSortDirection,
},
observability::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
@@ -64,6 +64,48 @@ const QUOTA_MAX_BYTES: usize = 128 * 1024;
const MAX_ADMIN_IMPORT_BATCH: usize = 200;
const ADMIN_USAGE_CHUNK_SIZE: usize = 200;
const ACCOUNT_USAGE_BY_WINDOWS_SQL: &str = "with requested_windows as (
select *
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
as requested(account_id, window_key, window_start, window_end)
)
select requested.account_id,
requested.window_key,
count(mr.id)::bigint as request_count,
count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count,
sum(mr.input_tokens)::bigint as input_tokens,
sum(mr.output_tokens)::bigint as output_tokens,
sum(mr.cached_tokens)::bigint as cached_tokens,
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
sum(mr.image_input_tokens)::bigint as image_input_tokens,
sum(mr.image_output_tokens)::bigint as image_output_tokens,
count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint
as image_request_count,
count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint
as image_request_failed_count,
coalesce(sum(coalesce(
mr.total_tokens,
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
)), 0)::bigint as total_tokens,
count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint
as provider_reported_count,
count(mr.id) filter (where mr.cost_source = 'calculated')::bigint
as calculated_count,
count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint
as unavailable_count,
max(mr.started_at) as last_used_at
from requested_windows requested
left join model_requests mr
on mr.provider_account_ref = requested.account_id
and mr.started_at >= requested.window_start
and mr.started_at < requested.window_end
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
group by requested.account_id, requested.window_key
order by requested.account_id, requested.window_key";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountAdminScope {
pub provider_kind: String,
@@ -898,6 +940,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
/// 三个 PostgreSQL adapter 都保持私有，调用方只能取得 [`AccountStore`] 暴露的领域能力。
#[derive(Clone)]
pub struct PgAdminAccountStore {
pool: PgPool,
accounts: PgProviderAccountRepository,
observability: PgObservabilityRepository,
control_plane: PgControlPlaneRepository,
@@ -907,6 +950,7 @@ impl PgAdminAccountStore {
#[must_use]
pub fn new(pool: PgPool) -> Self {
Self {
pool: pool.clone(),
accounts: PgProviderAccountRepository::new(pool.clone()),
observability: PgObservabilityRepository::new(pool.clone()),
control_plane: PgControlPlaneRepository::new(pool),
@@ -942,6 +986,61 @@ impl PgAdminAccountStore {
Ok(observations)
}
async fn usage_by_windows(
&self,
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
if windows.is_empty() {
return Ok(Vec::new());
}
let account_ids = windows
.iter()
.map(|window| window.account_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
validate_admin_account_ids(&account_ids)
.map_err(|error| admin_store_error(ENTITY, error))?;
for window in windows {
require_nonempty(ENTITY, "quota window key", &window.key)
.map_err(|error| admin_store_error(ENTITY, error))?;
ObservabilityRange::new(window.range.start, window.range.end)
.map_err(|error| admin_store_error(ENTITY, error))?;
}
let keys = windows
.iter()
.map(|window| window.key.clone())
.collect::<Vec<_>>();
let starts = windows
.iter()
.map(|window| window.range.start)
.collect::<Vec<_>>();
let ends = windows
.iter()
.map(|window| window.range.end)
.collect::<Vec<_>>();
let account_ids = windows
.iter()
.map(|window| window.account_id.clone())
.collect::<Vec<_>>();
sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL)
.bind(account_ids)
.bind(keys)
.bind(starts)
.bind(ends)
.fetch_all(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("load provider account quota window usage"),
)
})?
.iter()
.map(admin_account_usage_window)
.collect()
}
async fn required_scope(
&self,
account_id: &str,
@@ -1164,6 +1263,13 @@ impl AccountStore for PgAdminAccountStore {
.collect()
}
async fn load_account_usage_by_windows(
&self,
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
self.usage_by_windows(windows).await
}
async fn list_credentials(
&self,
provider_kind: &ProviderKind,
@@ -1622,11 +1728,25 @@ fn admin_account_summary(
.count(),
)
.unwrap_or(u64::MAX);
let attention = u64::try_from(
accounts
.iter()
.filter(|account| {
matches!(
admin_account_status(account, now),
AdminAccountStatus::Expired
| AdminAccountStatus::Disabled
| AdminAccountStatus::Banned
)
})
.count(),
)
.unwrap_or(u64::MAX);
AccountSummary {
total,
active,
quota_exhausted,
attention: total.saturating_sub(active),
attention,
}
}
@@ -1804,6 +1924,74 @@ fn admin_account_usage(usage: ProviderAccountUsageObservation) -> AdminStoreResu
})
}
fn admin_account_usage_window(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<AccountUsageWindowResult> {
let provider_reported_count = window_usage_count(row, "provider_reported_count")?;
let calculated_count = window_usage_count(row, "calculated_count")?;
let unavailable_count = window_usage_count(row, "unavailable_count")?;
Ok(AccountUsageWindowResult {
account_id: window_usage_value(row, "account_id")?,
key: window_usage_value(row, "window_key")?,
usage: AccountUsage {
account_id: window_usage_value(row, "account_id")?,
request_count: window_usage_count(row, "request_count")?,
success_count: window_usage_count(row, "success_count")?,
input_tokens: optional_window_usage_count(row, "input_tokens")?,
output_tokens: optional_window_usage_count(row, "output_tokens")?,
cached_tokens: optional_window_usage_count(row, "cached_tokens")?,
cache_write_tokens: optional_window_usage_count(row, "cache_write_tokens")?,
reasoning_tokens: optional_window_usage_count(row, "reasoning_tokens")?,
image_input_tokens: optional_window_usage_count(row, "image_input_tokens")?,
image_output_tokens: optional_window_usage_count(row, "image_output_tokens")?,
image_request_count: window_usage_count(row, "image_request_count")?,
image_request_failed_count: window_usage_count(row, "image_request_failed_count")?,
total_tokens: optional_window_usage_count(row, "total_tokens")?,
cost_coverage: AdminCostCoverage {
provider_reported_count,
calculated_count,
partial_count: 0,
unavailable_count,
not_billable_count: 0,
},
costs: Vec::new(),
last_used_at: window_usage_value(row, "last_used_at")?,
request_buckets: Vec::new(),
models: Vec::new(),
},
})
}
fn window_usage_count(row: &sqlx::postgres::PgRow, column: &'static str) -> AdminStoreResult<u64> {
let value = window_usage_value::<i64>(row, column)?;
u64::try_from(value).map_err(|_| invalid_window_usage(column))
}
fn optional_window_usage_count(
row: &sqlx::postgres::PgRow,
column: &'static str,
) -> AdminStoreResult<Option<u64>> {
window_usage_value::<Option<i64>>(row, column)?
.map(|value| u64::try_from(value).map_err(|_| invalid_window_usage(column)))
.transpose()
}
fn window_usage_value<T>(row: &sqlx::postgres::PgRow, column: &'static str) -> AdminStoreResult<T>
where
for<'row> T: sqlx::Decode<'row, Postgres> + sqlx::Type<Postgres>,
{
row.try_get(column)
.map_err(|_| invalid_window_usage(column))
}
fn invalid_window_usage(column: &'static str) -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
format!("invalid quota window usage column: {column}"),
)
}
fn admin_account_model_usage(
usage: ProviderAccountModelUsageObservation,
) -> AdminStoreResult<AccountModelUsage> {
backend/crates/gateway-store/tests/postgres/execution.rs
+49
View File
@@ -264,6 +264,7 @@ fn successful_core_finalization(id: &str) -> CoreModelRequestFinalization {
upstream_transport: Some("websocket".to_owned()),
http_version: Some("HTTP/2".to_owned()),
websocket_pool: None,
provider_metadata_json: None,
error: None,
provider_error_code: None,
retry_after_ms: None,
@@ -275,6 +276,54 @@ fn successful_core_finalization(id: &str) -> CoreModelRequestFinalization {
}
}
#[tokio::test]
async fn core_adapter_should_persist_only_object_provider_observation() {
let Some(database) = TestDatabase::create("execution_provider_observation").await else {
return;
};
seed_running_request(&database.pool, "req_provider_observation")
.await
.expect("seed model request");
let store = PgExecutionStore::new(database.pool.clone());
let mut finalization = successful_core_finalization("req_provider_observation");
finalization.provider_metadata_json = Some("{\"effectiveModel\":\"gpt-test\"}".to_owned());
ExecutionStore::finalize_model_request(&store, finalization)
.await
.expect("persist provider observation");
let persisted: serde_json::Value = sqlx::query_scalar(
"select provider_observation_json from model_requests where id = 'req_provider_observation'",
)
.fetch_one(&database.pool)
.await
.expect("load provider observation");
assert_eq!(persisted["effectiveModel"], "gpt-test");
database.close().await;
}
#[tokio::test]
async fn core_adapter_should_reject_non_object_provider_observation() {
let Some(database) = TestDatabase::create("execution_invalid_provider_observation").await
else {
return;
};
seed_running_request(&database.pool, "req_invalid_provider_observation")
.await
.expect("seed model request");
let store = PgExecutionStore::new(database.pool.clone());
let mut finalization = successful_core_finalization("req_invalid_provider_observation");
finalization.provider_metadata_json = Some("[]".to_owned());
assert!(
ExecutionStore::finalize_model_request(&store, finalization)
.await
.is_err()
);
database.close().await;
}
async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into model_requests (
backend/crates/gateway-store/tests/postgres/mod.rs
+1 -1
View File
@@ -132,7 +132,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() {
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (9, 1));
assert_eq!((first_table_count, migration_count), (8, 1));
}
#[test]
backend/crates/gateway-store/tests/postgres/observability.rs
+157 -52
View File
@@ -62,6 +62,48 @@ fn postgres_admin_observability_adapter_implements_terminal_port() {
assert_port::<PgAdminObservabilityStore>();
}
#[tokio::test]
async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
let Some(database) = TestDatabase::create("calculated_usage_billing_facts").await else {
return;
};
let now = Utc::now();
seed_observability_facts(&database.pool, now)
.await
.expect("seed observability facts");
seed_calculated_billing_facts(&database.pool, now)
.await
.expect("seed calculated billing facts");
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
.expect("observability range");
let repository = PgObservabilityRepository::new(database.pool.clone());
let facts = repository
.usage_calculated_billing_facts(range, UsageRecordFilter::default())
.await
.expect("calculated usage billing facts");
assert_eq!(facts.len(), 1);
assert_eq!(facts[0].provider_kind, "openai");
assert_eq!(facts[0].upstream_model_id, "gpt-5.5");
assert_eq!(facts[0].input_tokens, Some(800));
assert_eq!(facts[0].output_tokens, Some(200));
assert_eq!(facts[0].total.amount.as_str(), "1.25");
let store = PgAdminObservabilityStore::new(database.pool.clone(), None);
let facts = store
.usage_calculated_billing_facts(
admin_observability::TimeRange::new(range.start, range.end)
.expect("admin observability range"),
admin_observability::UsageFilter::default(),
)
.await
.expect("admin calculated usage billing facts");
assert_eq!(facts.len(), 1);
assert_eq!(facts[0].total.amount.as_str(), "1.25");
database.close().await;
}
#[tokio::test]
async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_details() {
let Some(database) = TestDatabase::create("admin_observability").await else {
@@ -81,17 +123,17 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.await
.expect("admin dashboard summary");
assert_eq!(dashboard.range, range);
assert_eq!(dashboard.requests.request_count, 2);
assert_eq!(dashboard.requests.request_count, 1);
assert_eq!(dashboard.requests.first_token_latency_sum_ms, 120);
assert_eq!(dashboard.requests.latency_sum_ms, 1_600);
assert_eq!(dashboard.requests.min_latency_ms, Some(700));
assert_eq!(dashboard.requests.latency_sum_ms, 900);
assert_eq!(dashboard.requests.min_latency_ms, Some(900));
assert_eq!(dashboard.requests.max_latency_ms, Some(900));
assert_eq!(dashboard.attempts.attempt_count, 3);
assert_eq!(dashboard.attempts.attempt_count, 4);
assert_eq!(dashboard.attempts.cost_coverage.provider_reported_count, 1);
assert_eq!(dashboard.attempts.cost_coverage.unavailable_count, 1);
assert_eq!(dashboard.attempts.cost_coverage.unavailable_count, 0);
assert_eq!(dashboard.attempts.costs[0].amount.as_str(), "1.25");
assert_eq!(dashboard.provider_accounts.total, 1);
assert_eq!(dashboard.account_usage[0].request_count, 2);
assert_eq!(dashboard.account_usage[0].request_count, 1);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24);
assert_eq!(
dashboard.account_usage[0]
@@ -99,7 +141,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.iter()
.map(|bucket| bucket.request_count)
.sum::<u64>(),
2,
1,
);
assert_eq!(dashboard.recent_requests.len(), 1);
assert_eq!(
@@ -124,7 +166,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.iter()
.map(|point| point.metrics.request_count)
.sum::<u64>(),
2,
1,
);
let trend = store
@@ -165,20 +207,9 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
})
.await
.expect("first usage page");
assert_eq!(first_page.total, 2);
assert_eq!(first_page.total, 1);
assert_eq!(first_page.items.len(), 1);
let second_page = store
.list_usage_records(admin_observability::UsageQuery {
range,
filter: admin_observability::UsageFilter::default(),
cursor: first_page.next_cursor,
page: admin_page(1),
page_size: PageSize::new(1).expect("page size"),
})
.await
.expect("second usage page");
assert_eq!(second_page.items.len(), 1);
assert_ne!(first_page.items[0].id, second_page.items[0].id);
assert!(first_page.next_cursor.is_none());
let filtered = store
.list_usage_records(admin_observability::UsageQuery {
@@ -230,21 +261,29 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
assert!(other_outcome.items.is_empty());
let detail = store
.usage_record_detail("req_observe_failed")
.usage_record_detail("req_observe_success")
.await
.expect("admin usage detail");
assert_eq!(
detail.request.outcome,
admin_observability::RequestOutcome::Failed
admin_observability::RequestOutcome::Succeeded
);
assert_eq!(detail.attempts.len(), 2);
assert_eq!(detail.attempts.len(), 1);
assert_eq!(
detail.attempts[0].outcome,
admin_observability::RequestOutcome::Failed
admin_observability::RequestOutcome::Succeeded
);
assert_eq!(
detail.attempts[1].outcome,
admin_observability::RequestOutcome::Failed
assert!(
store
.usage_record_detail("req_observe_failed")
.await
.is_err()
);
assert!(
store
.usage_record_detail("req_observe_uncommitted")
.await
.is_err()
);
let overview = store
@@ -352,10 +391,10 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.dashboard_summary(range)
.await
.expect("dashboard summary");
assert_eq!(dashboard.requests.request_count, 2);
assert_eq!(dashboard.requests.cache_eligible_request_count, 2);
assert_eq!(dashboard.requests.request_count, 1);
assert_eq!(dashboard.requests.cache_eligible_request_count, 1);
assert_eq!(dashboard.requests.cache_hit_request_count, 1);
assert_eq!(dashboard.requests.cache_hit_request_rate(), Some(0.5));
assert_eq!(dashboard.requests.cache_hit_request_rate(), Some(1.0));
assert_eq!(
dashboard
.requests
@@ -363,7 +402,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.p50_ms
.expect("latency p50")
.as_f64(),
800.0
900.0
);
assert_eq!(
dashboard
@@ -372,7 +411,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.p95_ms
.expect("latency p95")
.as_f64(),
890.0
900.0
);
assert_eq!(
dashboard
@@ -381,7 +420,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.p99_ms
.expect("latency p99")
.as_f64(),
898.0
900.0
);
assert_eq!(
dashboard
@@ -393,7 +432,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
120.0
);
assert_eq!(dashboard.provider_accounts.total, 1);
assert_eq!(dashboard.account_usage[0].request_count, 2);
assert_eq!(dashboard.account_usage[0].request_count, 1);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24);
assert_eq!(
dashboard.account_usage[0]
@@ -401,7 +440,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.iter()
.map(|bucket| bucket.request_count)
.sum::<u64>(),
2,
1,
);
assert_eq!(dashboard.recent_requests.len(), 1);
assert_eq!(dashboard.recent_requests[0].id, "req_observe_success");
@@ -411,7 +450,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.iter()
.map(|point| point.metrics.request_count)
.sum::<u64>(),
2
1
);
let account_usage = repository
@@ -423,11 +462,11 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
)
.await
.expect("provider account usage");
assert_eq!(account_usage[0].request_count, 2);
assert_eq!(account_usage[0].request_count, 1);
assert_eq!(account_usage[0].authentication_kind, "oauth");
assert_eq!(account_usage[0].models[0].request_count, 2);
assert_eq!(account_usage[0].models[0].request_count, 1);
assert_eq!(account_usage[0].cost_coverage.provider_reported_count, 1);
assert_eq!(account_usage[0].cost_coverage.unavailable_count, 1);
assert_eq!(account_usage[0].cost_coverage.unavailable_count, 0);
assert_eq!(account_usage[0].costs[0].amount.as_str(), "1.25");
assert_eq!(
account_usage[0]
@@ -435,7 +474,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.iter()
.map(|bucket| bucket.request_count)
.collect::<Vec<_>>(),
vec![2, 0],
vec![1, 0],
);
assert_eq!(
(
@@ -446,7 +485,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
account_usage[0].models[0].image_request_count,
account_usage[0].models[0].image_request_failed_count,
),
(Some(31), Some(9), 1, 1, 1, 1)
(Some(31), Some(9), 1, 0, 1, 0)
);
let usage_page = repository
@@ -462,7 +501,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
})
.await
.expect("usage records");
assert_eq!(usage_page.total, 2);
assert_eq!(usage_page.total, 1);
let successful_image = usage_page
.items
.iter()
@@ -489,20 +528,20 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
)
);
let detail = repository
.usage_record_detail("req_observe_failed")
.await
.expect("usage detail");
assert_eq!(detail.attempts.len(), 2);
assert_eq!(detail.attempts[0].source, "ops_event");
assert_eq!(detail.attempts[1].source, "model_request");
assert!(
repository
.usage_record_detail("req_observe_failed")
.await
.is_err()
);
let overview = repository
.usage_summary(range, UsageRecordFilter::default())
.await
.expect("usage summary");
assert_eq!(overview.attempts.attempt_count, 3);
assert_eq!(overview.attempts.attempt_count, 4);
assert_eq!(overview.attempts.failure_count, 2);
assert_eq!(overview.requests.request_count, 1);
let succeeded = repository
.usage_summary(
@@ -536,7 +575,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.await
.expect("usage diagnostics");
assert_eq!(diagnostics[0].name, "acct_observe");
assert_eq!(diagnostics[0].request_count, 2);
assert_eq!(diagnostics[0].request_count, 1);
assert_eq!(diagnostics[0].costs[0].amount.as_str(), "1.25");
let errors = repository
@@ -574,6 +613,27 @@ async fn seed_observability_facts(
.bind(now)
.execute(pool)
.await?;
sqlx::query(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport,
attempt_count, upstream_send_state, outcome, client_status_code,
input_tokens, output_tokens, total_tokens, cost_source, latency_ms,
started_at, deadline_at, completed_at
) values (
'req_observe_uncommitted', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
'acct_observe', 'upstream-model', 'http_sse',
1, 'sent', 'succeeded', 200,
900, 900, 1800, 'unavailable', 650,
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes'
)",
)
.bind(now)
.execute(pool)
.await?;
sqlx::query(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
@@ -643,3 +703,48 @@ async fn seed_observability_facts(
.await?;
Ok(())
}
async fn seed_calculated_billing_facts(
pool: &PgPool,
now: chrono::DateTime<Utc>,
) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens,
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at
) values (
'req_observe_calculated', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe', 'acct_observe', 'gpt-5.5',
'http_sse', 1, 'sent', $1 - interval '29 minutes', 'succeeded', 200,
800, 200, 0, 0, 1000, 'calculated', 1.25, 'USD',
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes'
)",
)
.bind(now)
.execute(pool)
.await?;
sqlx::query(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, outcome, client_status_code, input_tokens, output_tokens,
total_tokens, cost_source, cost_amount, cost_currency, started_at, deadline_at,
completed_at
) values (
'req_observe_calculated_uncommitted', 'key_observe', 1, 'openai', 'responses',
'/v1/responses', 'http_sse', 'public-model', 'openai', 'acct_observe',
'acct_observe', 'gpt-5.5', 'http_sse', 1, 'sent', 'succeeded', 200, 800, 200,
1000, 'calculated', 1.25, 'USD',
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes'
)",
)
.bind(now)
.execute(pool)
.await?;
Ok(())
}
backend/crates/gateway-store/tests/postgres/provider_accounts.rs
+104 -9
View File
@@ -5,8 +5,9 @@ use gateway_admin::{
model::{
MutationActor, MutationContext, PageSize, Revision as AdminRevision,
accounts::{
AccountListQuery, AccountSort, AccountSortField, AccountStatus, DeleteAccounts,
SetAccountEnabled as AdminSetAccountEnabled, SortDirection,
AccountListQuery, AccountSort, AccountSortField, AccountStatus,
AccountUsageWindowQuery, DeleteAccounts, SetAccountEnabled as AdminSetAccountEnabled,
SortDirection,
},
observability::TimeRange,
provider_credentials::{
@@ -128,7 +129,10 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
let mut attention = account("acct_attention", "user-attention");
attention.email = Some("attention@example.invalid".to_owned());
attention.availability = ProviderAccountAvailability::Invalid;
for account in [alpha, beta, charlie, attention] {
let mut quota_exhausted = account("acct_quota_exhausted", "user-quota-exhausted");
quota_exhausted.email = Some("quota-exhausted@example.invalid".to_owned());
quota_exhausted.availability = ProviderAccountAvailability::QuotaExhausted;
for account in [alpha, beta, charlie, attention, quota_exhausted] {
repository
.insert_provider_account(account)
.await
@@ -208,11 +212,11 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
.await
.expect("sort accounts by retained usage");
assert_eq!(usage_page.config_revision.get(), 1);
assert_eq!(usage_page.total, 4);
assert_eq!(usage_page.summary.total, 4);
assert_eq!(usage_page.total, 5);
assert_eq!(usage_page.summary.total, 5);
assert_eq!(usage_page.summary.active, 2);
assert_eq!(usage_page.summary.quota_exhausted, 0);
assert_eq!(usage_page.summary.attention, 2);
assert_eq!(usage_page.summary.quota_exhausted, 1);
assert_eq!(usage_page.summary.attention, 1);
assert_eq!(
usage_page
.items
@@ -427,6 +431,96 @@ async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs(
database.close().await;
}
#[tokio::test]
async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back_to_components() {
let Some(database) = TestDatabase::create("provider_account_quota_window_usage").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
repository
.insert_provider_account(account("acct_quota_window", "user-quota-window"))
.await
.expect("insert quota window account");
let now = Utc::now();
seed_model_request(
&database.pool,
ModelRequestSeed {
request_id: "req_quota_window_recent",
account_id: "acct_quota_window",
provider_kind: "openai",
model: "gpt-window",
total_tokens: 10,
cost_amount: "0",
started_at: now - TimeDelta::minutes(30),
},
)
.await
.expect("seed recent quota window request");
seed_model_request(
&database.pool,
ModelRequestSeed {
request_id: "req_quota_window_older",
account_id: "acct_quota_window",
provider_kind: "openai",
model: "gpt-window",
total_tokens: 20,
cost_amount: "0",
started_at: now - TimeDelta::hours(6),
},
)
.await
.expect("seed older quota window request");
sqlx::query(
"update model_requests
set input_tokens = case id
when 'req_quota_window_recent' then 7
when 'req_quota_window_older' then 11
end,
output_tokens = case id
when 'req_quota_window_recent' then 3
when 'req_quota_window_older' then 9
end,
total_tokens = case id
when 'req_quota_window_recent' then 700
when 'req_quota_window_older' then null
end
where id in ('req_quota_window_recent', 'req_quota_window_older')",
)
.execute(&database.pool)
.await
.expect("separate provider total from fallback components");
let mut usage = PgAdminAccountStore::new(database.pool.clone())
.load_account_usage_by_windows(&[
AccountUsageWindowQuery {
account_id: "acct_quota_window".to_owned(),
key: "short".to_owned(),
range: TimeRange {
start: now - TimeDelta::hours(1),
end: now + TimeDelta::minutes(1),
},
},
AccountUsageWindowQuery {
account_id: "acct_quota_window".to_owned(),
key: "long".to_owned(),
range: TimeRange {
start: now - TimeDelta::hours(7),
end: now + TimeDelta::minutes(1),
},
},
])
.await
.expect("load quota window usage");
usage.sort_by(|left, right| left.key.cmp(&right.key));
assert_eq!(usage[0].key, "long");
assert_eq!(usage[0].usage.total_tokens, Some(720));
assert_eq!(usage[1].key, "short");
assert_eq!(usage[1].usage.total_tokens, Some(700));
database.close().await;
}
#[tokio::test]
async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
let Some(database) = TestDatabase::create("provider_account_terminal_mutation").await else {
@@ -1014,14 +1108,15 @@ async fn seed_model_request(
client_transport, requested_model_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, outcome, client_status_code, upstream_status_code,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
upstream_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens,
total_tokens, cost_source, cost_amount, cost_currency,
started_at, deadline_at, completed_at
) values (
$1, 'key-provider-account-test', 1, 'openai', 'responses', '/v1/responses',
'http_sse', $4, $3, $2, $2, $4, 'http_sse', 1,
'sent', 'succeeded', 200, 200, $5, 0, 0, 0, 0,
'sent', $7 + interval '1 second', 'succeeded', 200, 200, $5, 0, 0, 0, 0,
$5, 'provider_reported', $6::numeric, 'USD', $7,
$7 + interval '5 minutes', $7 + interval '1 second'
)",
backend/crates/providers/openai/src/config.rs
+1 -1
View File
@@ -159,7 +159,7 @@ impl Default for CodexWebSocketPoolSettings {
max_per_account: 8,
max_total: 64,
max_connecting: 8,
initial_event_timeout_ms: 800,
initial_event_timeout_ms: 20_000,
}
}
}
backend/crates/providers/openai/src/provider.rs
+431 -45
View File
@@ -3,7 +3,7 @@
use std::collections::BTreeSet;
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use futures::{StreamExt, future::BoxFuture};
@@ -22,8 +22,8 @@ use gateway_core::error::{
SafeUpstreamValue,
};
use gateway_core::event::{
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation,
ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind,
GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseMetadata,
ProviderResponseObservation, ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind,
};
use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProviderSessionState};
use gateway_core::provider_ports::ProviderSessionAffinityKey;
@@ -36,7 +36,7 @@ use gateway_core::task::{
use reqwest::Client;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value, json};
use url::Url;
use crate::credential::{
@@ -57,7 +57,9 @@ use crate::transport::profile::{
APPCAST_POLL_INTERVAL, CodexCliReleaseService, CodexDesktopReleaseService,
CodexWireProfileState,
};
use crate::transport::protocol::responses::{CodexResponsesRequest, PreviousResponseScope};
use crate::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
};
use crate::transport::request::{
CodexRequestEncodeError, encode_generate_request, sanitize_cross_account_item,
scope_request_to_account,
@@ -66,8 +68,8 @@ use crate::transport::session::CodexSessionIdentity;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics,
CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
CodexClientError, CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
};
const PROVIDER_NAME: &str = "openai";
@@ -779,31 +781,13 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
});
capture.turn_state = response.turn_state.clone().or(capture.turn_state.clone());
}
let observation = match codex_response_observation(
response.transport,
&response.diagnostics,
&response.response_metadata,
&response.transport_metrics,
response.websocket_pool_decision,
) {
Ok(observation) => observation,
Err(error) => {
let failure = MappedProviderFailure::plain(error);
apply_failure(
&client,
&selector,
&quota,
&active_account,
&response_origin,
cyber_policy_scope.as_ref(),
&failure,
)
.await;
Err(failure.error)?;
return;
}
};
yield ProviderEvent::observation(observation);
let mut observation_state = OpenAiResponseObservationState::from_backend_response(
&response,
&request,
active_account.id().as_str(),
context.attempt_index().get(),
);
yield ProviderEvent::observation(observation_state.observation()?);
synchronize_passive_quota(&quota, &active_account, &response.rate_limit_headers).await;
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(etag)
@@ -861,6 +845,13 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
let next = match next {
Ok(next) => next,
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
synchronize_passive_quota(&quota, &active_account, &updates).await;
if observation_state.merge_rate_limit_headers(&updates) {
yield ProviderEvent::observation(observation_state.observation()?);
}
}
apply_failure(
&client,
&selector,
@@ -876,16 +867,21 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
};
let Some(chunk) = next else { break; };
if let Some(updates) = rate_limit_updates.as_ref() {
let updates = std::mem::take(&mut *updates.lock().await);
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
}
observation_state.merge_rate_limit_headers(&updates)
};
if let (Some(capture), Some(updates)) =
(session_capture.as_mut(), turn_state_updates.as_ref())
&& let Some(turn_state) = updates.lock().await.clone()
{
capture.turn_state = Some(turn_state);
}
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, context.timing_started_at());
let (mut events, terminal_failure) = match decoder.push(&chunk) {
CodexCanonicalOutcome::Events(events) => (events, None),
CodexCanonicalOutcome::Failed(failure) => {
@@ -900,6 +896,18 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
(events, Some(failure))
}
};
let timing_changed = first_event_changed
|| observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
let completed = events
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if rate_limits_changed || timing_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
@@ -931,10 +939,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
}
return;
}
let completed = events
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
for event in events {
yield event;
}
@@ -966,6 +970,15 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
(events, Some(failure))
}
};
let timing_changed = observation_state
.observe_timing_signals(decoder.take_timing_signals(), context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
let rate_limits_changed = if updates.is_empty() {
false
} else {
synchronize_passive_quota(&quota, &active_account, &updates).await;
observation_state.merge_rate_limit_headers(&updates)
};
if let Some(failure) = terminal_failure.as_ref() {
apply_failure(
&client,
@@ -1007,6 +1020,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
.iter()
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if timing_changed || rate_limits_changed || terminal_changed {
yield ProviderEvent::observation(observation_state.observation()?);
}
for event in events {
yield event;
}
@@ -1025,23 +1043,258 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
})
}
/// OpenAI Responses 的观测状态完全归 Provider 所有。
///
/// 每次原始 SSE/WS 事件推进状态后重新生成不可变 observation，Core 只负责携带、
/// 持久化和展示这份安全快照，不解释任何 OpenAI 协议字段。
struct OpenAiResponseObservationState {
transport: CodexBackendTransport,
diagnostics: CodexUpstreamDiagnostics,
response_metadata: CodexResponseMetadata,
metrics: CodexTransportMetrics,
websocket_pool_decision: Option<crate::transport::WebSocketPoolDecision>,
request_summary: Value,
requested_model: String,
stream: bool,
compact: bool,
service_tier: Option<String>,
rate_limit_headers: Vec<(String, String)>,
account_id: String,
attempt_index: u32,
timings: ProviderResponseTimings,
terminal: Option<OpenAiResponseTerminal>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum OpenAiResponseTerminal {
Completed,
Incomplete,
}
impl OpenAiResponseObservationState {
fn from_backend_response(
response: &CodexBackendStreamingResponse,
request: &CodexResponsesRequest,
account_id: &str,
attempt_index: u32,
) -> Self {
let semantics = request.semantics();
Self {
transport: response.transport,
diagnostics: response.diagnostics.clone(),
response_metadata: response.response_metadata.clone(),
metrics: response.transport_metrics.clone(),
websocket_pool_decision: response.websocket_pool_decision,
request_summary: openai_response_request_summary(request, response.transport),
requested_model: request.model().to_owned(),
stream: request.stream(),
compact: semantics.compact,
service_tier: request.service_tier().map(str::to_owned),
rate_limit_headers: safe_observation_headers(&response.rate_limit_headers),
account_id: account_id.to_owned(),
attempt_index,
timings: openai_response_timings(
&response.transport_metrics,
&response.response_metadata,
),
terminal: None,
}
}
fn observation(&self) -> Result<ProviderResponseObservation, ProviderError> {
let mut observation = codex_response_observation(
self.transport,
&self.diagnostics,
&self.response_metadata,
&self.metrics,
self.websocket_pool_decision,
self.timings,
)?;
if let Some(metadata) = self.provider_metadata() {
observation = observation.with_provider_metadata(metadata);
}
Ok(observation)
}
fn observe_stream_chunk(&mut self, chunk: &[u8], started_at: Instant) -> bool {
if chunk.is_empty() {
return false;
}
insert_first_timing(&mut self.timings.first_event_ms, started_at)
}
fn observe_timing_signals(
&mut self,
signals: ResponseEventSignals,
started_at: Instant,
) -> bool {
let mut changed = false;
if signals.semantic_output {
changed |= insert_first_timing(&mut self.timings.first_token_ms, started_at);
}
if signals.reasoning_output {
changed |= insert_first_timing(&mut self.timings.first_reasoning_ms, started_at);
}
if signals.text_output {
changed |= insert_first_timing(&mut self.timings.first_text_ms, started_at);
}
changed
}
fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
let mut changed = false;
for (name, value) in safe_observation_headers(updates) {
if let Some(existing) = self
.rate_limit_headers
.iter_mut()
.find(|(existing_name, _)| existing_name == &name)
{
if existing.1 != value {
existing.1 = value;
changed = true;
}
} else {
self.rate_limit_headers.push((name, value));
changed = true;
}
}
changed
}
fn mark_completed(&mut self, incomplete: bool) -> bool {
let terminal = if incomplete {
OpenAiResponseTerminal::Incomplete
} else {
OpenAiResponseTerminal::Completed
};
if self.terminal == Some(terminal) {
return false;
}
self.terminal = Some(terminal);
true
}
fn provider_metadata(&self) -> Option<ProviderResponseMetadata> {
let mut metadata = Map::new();
let effective_model = self
.response_metadata
.effective_model
.as_deref()
.filter(|model| !model.is_empty())
.unwrap_or(&self.requested_model);
if !effective_model.is_empty() {
metadata.insert(
"effectiveModel".to_owned(),
Value::String(effective_model.to_owned()),
);
}
metadata.insert(
"modelsEtag".to_owned(),
self.response_metadata
.models_etag
.as_ref()
.map_or(Value::Null, |etag| Value::String(etag.clone())),
);
metadata.insert(
"reasoningIncluded".to_owned(),
Value::Bool(self.response_metadata.reasoning_included),
);
metadata.insert("stream".to_owned(), Value::Bool(self.stream));
metadata.insert("compact".to_owned(), Value::Bool(self.compact));
metadata.insert("attemptIndex".to_owned(), json!(self.attempt_index));
metadata.insert(
"attemptAccountId".to_owned(),
Value::String(self.account_id.clone()),
);
metadata.insert(
"rateLimitHeaders".to_owned(),
json!(self.rate_limit_headers),
);
metadata.insert(
"upstreamTraceHeaders".to_owned(),
json!(safe_observation_headers(&self.diagnostics.trace_headers)),
);
metadata.insert("requestSummary".to_owned(), self.request_summary.clone());
if let Some(service_tier) = &self.service_tier {
metadata.insert(
"serviceTier".to_owned(),
Value::String(service_tier.clone()),
);
}
if let Some(status) = self.diagnostics.status_code {
metadata.insert("upstreamStatus".to_owned(), json!(status));
}
if let Some(decision) = self.metrics.decision {
metadata.insert(
"transportDecision".to_owned(),
Value::String(decision.as_str().to_owned()),
);
}
if let Some(decision) = self.websocket_pool_decision {
metadata.insert(
"websocketPool".to_owned(),
json!({ "kind": decision.kind() }),
);
}
if let Some(version) = self.metrics.http_version.as_deref() {
metadata.insert("httpVersion".to_owned(), Value::String(version.to_owned()));
}
if let Some((_, cf_ray)) = self
.diagnostics
.trace_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("cf-ray"))
&& let Some((_, cf_ray)) = safe_observation_header("cf-ray", cf_ray)
{
metadata.insert("cfRay".to_owned(), Value::String(cf_ray));
}
insert_optional_metadata_millis(
&mut metadata,
"transportDecisionWaitMs",
self.timings.transport_decision_wait_ms,
);
insert_optional_metadata_millis(&mut metadata, "wsConnectMs", self.timings.connect_ms);
insert_optional_metadata_millis(
&mut metadata,
"upstreamHeadersMs",
self.timings.headers_ms,
);
insert_optional_metadata_millis(&mut metadata, "firstEventMs", self.timings.first_event_ms);
insert_optional_metadata_millis(
&mut metadata,
"firstReasoningMs",
self.timings.first_reasoning_ms,
);
insert_optional_metadata_millis(&mut metadata, "firstTextMs", self.timings.first_text_ms);
insert_optional_metadata_millis(&mut metadata, "firstTokenMs", self.timings.first_token_ms);
insert_optional_metadata_millis(
&mut metadata,
"openaiProcessingMs",
self.timings.provider_processing_ms,
);
if let Some(terminal) = self.terminal {
let incomplete = terminal == OpenAiResponseTerminal::Incomplete;
metadata.insert("completed".to_owned(), Value::Bool(!incomplete));
metadata.insert("incomplete".to_owned(), Value::Bool(incomplete));
metadata.insert("eventStatusCode".to_owned(), json!(200_u16));
}
ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?)
}
}
fn codex_response_observation(
transport: CodexBackendTransport,
diagnostics: &CodexUpstreamDiagnostics,
response_metadata: &CodexResponseMetadata,
metrics: &CodexTransportMetrics,
websocket_pool_decision: Option<crate::transport::WebSocketPoolDecision>,
timings: ProviderResponseTimings,
) -> Result<ProviderResponseObservation, ProviderError> {
let mut observation = ProviderResponseObservation::new(
UpstreamTransport::new(actual_transport_name(transport))
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?,
)
.with_timings(ProviderResponseTimings {
transport_decision_wait_ms: nonnegative_millis(metrics.transport_decision_wait_ms),
connect_ms: nonnegative_millis(metrics.ws_connect_ms),
headers_ms: nonnegative_millis(metrics.upstream_headers_ms),
first_event_ms: nonnegative_millis(metrics.first_event_ms),
});
.with_timings(timings);
if let Some(version) = metrics
.http_version
.as_deref()
@@ -1097,6 +1350,7 @@ fn codex_error_observation(error: &CodexClientError) -> Option<ProviderResponseO
&CodexResponseMetadata::default(),
transport_metrics,
None,
openai_response_timings(transport_metrics, &CodexResponseMetadata::default()),
)
.ok()?
.with_status_code(status.as_u16());
@@ -1117,6 +1371,138 @@ fn codex_error_observation(error: &CodexClientError) -> Option<ProviderResponseO
Some(observation)
}
fn openai_response_timings(
metrics: &CodexTransportMetrics,
response_metadata: &CodexResponseMetadata,
) -> ProviderResponseTimings {
ProviderResponseTimings {
transport_decision_wait_ms: nonnegative_millis(metrics.transport_decision_wait_ms),
connect_ms: nonnegative_millis(metrics.ws_connect_ms),
headers_ms: nonnegative_millis(metrics.upstream_headers_ms),
first_event_ms: nonnegative_millis(metrics.first_event_ms),
provider_processing_ms: openai_processing_ms(response_metadata),
..ProviderResponseTimings::default()
}
}
fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option<u64> {
response_metadata
.client_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms"))
.and_then(|(_, value)| value.parse::<u64>().ok())
}
fn openai_response_request_summary(
request: &CodexResponsesRequest,
transport: CodexBackendTransport,
) -> Value {
let body = request.body();
let input = body.get("input");
let tools = body.get("tools");
let semantics = request.semantics();
json!({
"model": request.model(),
"stream": request.stream(),
"store": request.store(),
"compact": semantics.compact,
"requestKind": semantics.request_kind,
"subagentKind": semantics.subagent_kind,
"reasoningPreset": semantics.reasoning_preset,
"transport": actual_transport_name(transport),
"inputType": json_value_kind(input),
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
"toolsType": json_value_kind(tools),
"toolsCount": tools.and_then(Value::as_array).map(Vec::len),
"topLevelFields": body.keys().cloned().collect::<Vec<_>>(),
"previousResponseIdPresent": request.previous_response_id().is_some(),
"serviceTier": request.service_tier(),
"localTransport": {
"useWebsocket": request.use_websocket,
"forceHttpSse": request.force_http_sse,
},
})
}
const fn json_value_kind(value: Option<&Value>) -> &'static str {
match value {
None => "missing",
Some(Value::Null) => "null",
Some(Value::Bool(_)) => "boolean",
Some(Value::Number(_)) => "number",
Some(Value::String(_)) => "string",
Some(Value::Array(_)) => "array",
Some(Value::Object(_)) => "object",
}
}
fn safe_observation_headers(headers: &[(String, String)]) -> Vec<(String, String)> {
headers
.iter()
.filter_map(|(name, value)| safe_observation_header(name, value))
.collect()
}
fn safe_observation_header(name: &str, value: &str) -> Option<(String, String)> {
let name = name.trim().to_ascii_lowercase();
let value = value.trim();
let safe_value = SafeUpstreamValue::new(value.to_owned()).ok()?;
ProviderResponseHeader::new(name.clone(), safe_value)?;
Some((name, value.to_owned()))
}
fn insert_first_timing(target: &mut Option<u64>, started_at: Instant) -> bool {
if target.is_some() {
return false;
}
let elapsed = u64::try_from(started_at.elapsed().as_millis())
.unwrap_or(u64::MAX)
.max(1);
*target = Some(elapsed);
true
}
fn insert_optional_metadata_millis(
metadata: &mut Map<String, Value>,
name: &str,
value: Option<u64>,
) {
if let Some(value) = value {
metadata.insert(name.to_owned(), json!(value));
}
}
async fn take_rate_limit_updates(
updates: Option<&CodexRateLimitHeaderUpdates>,
) -> Vec<(String, String)> {
let Some(updates) = updates else {
return Vec::new();
};
std::mem::take(&mut *updates.lock().await)
}
fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
events
.iter()
.find_map(|event| {
let wire = event.wire_event()?;
let event_type = wire
.event_type()
.or_else(|| wire.data().get("type").and_then(Value::as_str));
match event_type {
Some("response.incomplete") => Some(true),
Some("response.completed") => Some(
wire.data()
.pointer("/response/status")
.and_then(Value::as_str)
== Some("incomplete"),
),
_ => None,
}
})
.unwrap_or(false)
}
async fn synchronize_passive_quota(
quota: &CodexCredentialQuotaService,
account: &ProviderAccount,
backend/crates/providers/openai/src/transport/canonical.rs
+26 -3
View File
@@ -18,7 +18,9 @@ use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder, SseFrame, sse_fra
use serde_json::Value;
use thiserror::Error;
use super::protocol::responses::{ResponsesSseFailure, response_event_signals};
use super::protocol::responses::{
ResponseEventSignals, ResponsesSseFailure, response_event_signals,
};
use super::usage::openai_billing_breakdown;
const CONTENTS_PER_OUTPUT: u32 = 1_024;
@@ -39,6 +41,7 @@ pub struct CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet<u32>,
usage_emitted: bool,
semantic_output_seen: bool,
timing_signals: ResponseEventSignals,
raw_sse_passthrough: bool,
}
@@ -132,6 +135,7 @@ impl CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet::new(),
usage_emitted: false,
semantic_output_seen: false,
timing_signals: ResponseEventSignals::default(),
raw_sse_passthrough: false,
}
}
@@ -146,6 +150,7 @@ impl CodexCanonicalDecoder {
}
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
let frames = match self.decoder.push_frames(chunk) {
Ok(frames) => frames,
@@ -171,6 +176,7 @@ impl CodexCanonicalDecoder {
}
pub fn finish(&mut self) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
let frames = match self.decoder.finish_frames() {
Ok(frames) => frames,
@@ -195,6 +201,15 @@ impl CodexCanonicalDecoder {
self.decode(events)
}
/// 取走最近一次解码中由原始 Responses 事件观察到的计时语义。
///
/// 这份事实独立于 canonical 投影：未知的未来事件仍可保留 wire 透明转发，
/// 同时让 Provider 正确记录首个可消费输出的时延。
#[must_use]
pub fn take_timing_signals(&mut self) -> ResponseEventSignals {
std::mem::take(&mut self.timing_signals)
}
fn decode(&mut self, events: Vec<SseEvent>) -> CodexCanonicalOutcome {
let mut output = Vec::new();
for event in events {
@@ -254,6 +269,8 @@ impl CodexCanonicalDecoder {
.event
.as_deref()
.or_else(|| value.get("type").and_then(Value::as_str));
let signals = response_event_signals(event_type, &value);
self.merge_timing_signals(signals);
if matches!(event_type, Some("response.failed" | "error")) {
let failure = ResponsesSseFailure::from_event(event_type.unwrap_or_default(), &value);
if let Ok(wire) = Self::wire_for_event(event, value, raw_sse_frame) {
@@ -265,7 +282,6 @@ impl CodexCanonicalDecoder {
if let Some(event_type) = event_type {
let _ = self.decode_event(event_type, &value, &mut canonical);
}
let semantic_output = response_event_signals(event_type, &value).semantic_output;
let Ok(wire) = Self::wire_for_event(event, value, raw_sse_frame) else {
return Ok(());
};
@@ -274,10 +290,17 @@ impl CodexCanonicalDecoder {
} else {
ProviderEvent::canonical_with_wire(canonical, wire)
});
self.semantic_output_seen |= semantic_output;
self.semantic_output_seen |= signals.semantic_output;
Ok(())
}
fn merge_timing_signals(&mut self, signals: ResponseEventSignals) {
self.timing_signals.protocol_progress |= signals.protocol_progress;
self.timing_signals.semantic_output |= signals.semantic_output;
self.timing_signals.reasoning_output |= signals.reasoning_output;
self.timing_signals.text_output |= signals.text_output;
}
fn wire_for_event(
event: SseEvent,
value: Value,
backend/crates/providers/openai/tests/config.rs
+1 -1
View File
@@ -79,7 +79,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
8,
64,
8,
800,
20_000,
5,
true,
true,
backend/crates/providers/openai/tests/transport/canonical.rs
+40
View File
@@ -172,6 +172,46 @@ fn decoder_should_restore_done_only_reasoning_and_text_as_canonical_facts() {
);
}
#[test]
fn decoder_timing_signals_should_count_tool_arguments_as_first_token() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_tool_timing\",\"model\":\"gpt-test\"}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"call_id\":\"call_tool_timing\"}}\n\n",
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"delta\":\"{\\\"path\\\":\\\"a\\\"}\"}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let _ = decoder.push(body.as_bytes()).expect("tool argument frame");
let signals = decoder.take_timing_signals();
assert!(signals.semantic_output);
assert!(!signals.reasoning_output);
assert!(!signals.text_output);
}
#[test]
fn decoder_timing_signals_should_distinguish_reasoning_and_text_output() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_output_timing\",\"model\":\"gpt-test\"}}\n\n",
"event: response.reasoning_summary_text.delta\n",
"data: {\"type\":\"response.reasoning_summary_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"plan\"}\n\n",
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":1,\"content_index\":0,\"delta\":\"answer\"}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let _ = decoder.push(body.as_bytes()).expect("output frames");
let signals = decoder.take_timing_signals();
assert!(signals.semantic_output);
assert!(signals.reasoning_output);
assert!(signals.text_output);
}
#[test]
fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() {
let body = concat!(
backend/crates/providers/xai/src/credential/import.rs
+132 -126
View File
@@ -1,15 +1,17 @@
//! 已有 Grok Build OAuth token 的安全归一化边界。
use std::collections::{HashMap, HashSet};
use std::collections::HashSet;
use std::fmt;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use serde::Deserialize;
use serde_json::Value;
use serde_json::{Map, Value};
use url::Url;
use crate::credential::token::UnverifiedTokenSet;
use crate::{FailureClass, GROK_CLI_BASE_URL, GrokOAuthConfig, OAuthError, SecretValue};
use crate::{
FailureClass, GROK_CLI_BASE_URL, GrokOAuthConfig, OAuthError, OFFICIAL_CLIENT_ID,
OFFICIAL_SCOPES, SecretValue,
};
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
@@ -143,88 +145,37 @@ impl fmt::Debug for GrokOAuthImportEntry {
}
}
/// Provider-owned xAI OAuth 账号文档。
/// Provider-owned xAI OAuth 账号导入。
pub struct GrokOAuthImportDocument {
entries: Vec<GrokOAuthImportEntry>,
}
impl GrokOAuthImportDocument {
/// 解析完整 OAuth 账号文档；未知字段、代理、API Key 或非 xAI OAuth 项均拒绝。
/// 从外部 JSON 提取 xAI OAuth 认证字段。
///
/// 来源包装格式、代理、并发和其他展示 metadata 不参与认证，也不会影响导入。
/// 实际 token 仍须通过官方 refresh/user-info 验证；API Key 不能混入 OAuth 条目。
pub fn parse_json(document: &[u8]) -> Result<Self, GrokOAuthImportError> {
if document.is_empty() || document.len() > MAX_IMPORT_DOCUMENT_BYTES {
return Err(GrokOAuthImportError::InvalidField("document"));
}
let wire: OAuthAccountDocumentWire = serde_json::from_slice(document)
let wire: Value = serde_json::from_slice(document)
.map_err(|_| GrokOAuthImportError::InvalidField("document"))?;
if wire
.kind
.as_deref()
.is_some_and(|kind| kind != "external-account-bundle")
|| wire.version.is_some_and(|version| version != 1)
|| !wire.proxies.is_empty()
|| wire.accounts.is_empty()
|| wire.accounts.len() > MAX_IMPORT_ACCOUNTS
{
let accounts = import_accounts(&wire)?;
if accounts.is_empty() || accounts.len() > MAX_IMPORT_ACCOUNTS {
return Err(GrokOAuthImportError::InvalidField("document"));
}
let mut names = HashSet::new();
let mut entries = Vec::with_capacity(wire.accounts.len());
for account in wire.accounts {
if account.credentials.extra_fields.keys().any(|field| {
matches!(
field.as_str(),
"api_key" | "apiKey" | "xai_api_key" | "xaiApiKey"
)
}) {
return Err(GrokOAuthImportError::InvalidField("account"));
let now = Utc::now();
let mut entries = Vec::with_capacity(accounts.len());
for (index, account) in accounts.into_iter().enumerate() {
if let Some(entry) = parse_account_entry(account, index, now)? {
entries.push(entry);
}
if account.platform != "grok"
|| account.kind != "oauth"
|| account.concurrency == 0
|| account.priority == 0
|| !valid_display_value(&account.name, MAX_ACCOUNT_NAME_BYTES)
|| !names.insert(account.name.clone())
{
return Err(GrokOAuthImportError::InvalidField("account"));
}
let email = match (account.credentials.email, account.extra.email) {
(Some(left), Some(right)) if left != right => {
return Err(GrokOAuthImportError::InvalidField("email"));
}
(Some(email), _) | (_, Some(email)) => Some(email),
(None, None) => None,
};
if email
.as_deref()
.is_some_and(|value| !valid_display_value(value, MAX_EMAIL_BYTES))
{
return Err(GrokOAuthImportError::InvalidField("email"));
}
let tokens = match account.credentials.id_token {
Some(id_token) => GrokOAuthImportTokens::new(
SecretValue::new(account.credentials.access_token),
SecretValue::new(account.credentials.refresh_token),
SecretValue::new(id_token),
),
None => GrokOAuthImportTokens::without_id_token(
SecretValue::new(account.credentials.access_token),
SecretValue::new(account.credentials.refresh_token),
),
};
let metadata = GrokOAuthImportMetadata::new(
account.credentials.token_type,
account.credentials.client_id,
account.credentials.scope,
account.credentials.base_url,
wire.exported_at,
account.credentials.expires_at,
);
entries.push(GrokOAuthImportEntry {
name: account.name,
email,
candidate: GrokOAuthImportCandidate::new(tokens, metadata),
});
}
if entries.is_empty() {
return Err(GrokOAuthImportError::InvalidField("document"));
}
Ok(Self { entries })
}
@@ -244,66 +195,121 @@ impl fmt::Debug for GrokOAuthImportDocument {
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct OAuthAccountDocumentWire {
#[serde(default, rename = "type")]
kind: Option<String>,
#[serde(default)]
version: Option<u32>,
exported_at: DateTime<Utc>,
accounts: Vec<OAuthAccountWire>,
proxies: Vec<Value>,
#[serde(default, rename = "skipped_shadows")]
_skipped_shadows: u64,
fn import_accounts(document: &Value) -> Result<Vec<&Value>, GrokOAuthImportError> {
match document {
Value::Array(accounts) => Ok(accounts.iter().collect()),
Value::Object(object) => match object.get("accounts") {
Some(Value::Array(accounts)) => Ok(accounts.iter().collect()),
Some(_) => Err(GrokOAuthImportError::InvalidField("accounts")),
None => Ok(vec![document]),
},
_ => Err(GrokOAuthImportError::InvalidField("document")),
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct OAuthAccountWire {
name: String,
#[serde(default, rename = "notes")]
_notes: Option<String>,
platform: String,
#[serde(rename = "type")]
kind: String,
credentials: OAuthCredentialsWire,
concurrency: u32,
priority: u32,
#[serde(default, rename = "rate_multiplier")]
_rate_multiplier: Option<Value>,
#[serde(default, rename = "expires_at")]
_expires_at: Option<i64>,
#[serde(default, rename = "auto_pause_on_expired")]
_auto_pause_on_expired: Option<bool>,
#[serde(default, rename = "proxy_key")]
_proxy_key: Option<String>,
#[serde(default)]
extra: OAuthAccountExtraWire,
fn parse_account_entry(
account: &Value,
index: usize,
now: DateTime<Utc>,
) -> Result<Option<GrokOAuthImportEntry>, GrokOAuthImportError> {
let account = account
.as_object()
.ok_or(GrokOAuthImportError::InvalidField("account"))?;
let credentials = account
.get("credentials")
.and_then(Value::as_object)
.unwrap_or(account);
let (access_token, refresh_token) = match (
credential_value(credentials, &["access_token", "accessToken", "at"]),
credential_value(credentials, &["refresh_token", "refreshToken", "rt"]),
) {
(Some(access_token), Some(refresh_token)) => (access_token, refresh_token),
(None, None) => return Ok(None),
(None, Some(_)) => return Err(GrokOAuthImportError::InvalidField("access_token")),
(Some(_), None) => return Err(GrokOAuthImportError::InvalidField("refresh_token")),
};
if ["api_key", "apiKey", "xai_api_key", "xaiApiKey"]
.iter()
.any(|field| credentials.contains_key(*field))
{
return Err(GrokOAuthImportError::InvalidField("account"));
}
let email = [
credential_value(credentials, &["email"]),
display_value(account, &["email"]),
account
.get("extra")
.and_then(Value::as_object)
.and_then(|extra| display_value(extra, &["email"])),
]
.into_iter()
.flatten()
.find(|value| valid_display_value(value, MAX_EMAIL_BYTES))
.map(str::to_owned);
let name = display_value(account, &["name", "label"])
.filter(|value| valid_display_value(value, MAX_ACCOUNT_NAME_BYTES))
.map(str::to_owned)
.or_else(|| email.clone())
.unwrap_or_else(|| format!("xAI OAuth account {}", index + 1));
let expires_at = credential_value(
credentials,
&[
"expires_at",
"expiresAt",
"access_token_expires_at",
"accessTokenExpiresAt",
],
)
.and_then(parse_source_expiry)
.filter(|expires_at| *expires_at > now && *expires_at - now <= MAX_DECLARED_LIFETIME)
.unwrap_or_else(|| now + ChronoDuration::seconds(1));
let tokens = match credential_value(credentials, &["id_token", "idToken"]) {
Some(id_token) => GrokOAuthImportTokens::new(
SecretValue::new(access_token.to_owned()),
SecretValue::new(refresh_token.to_owned()),
SecretValue::new(id_token.to_owned()),
),
None => GrokOAuthImportTokens::without_id_token(
SecretValue::new(access_token.to_owned()),
SecretValue::new(refresh_token.to_owned()),
),
};
let metadata = GrokOAuthImportMetadata::new(
"Bearer".to_owned(),
OFFICIAL_CLIENT_ID.to_owned(),
OFFICIAL_SCOPES.join(" "),
GROK_CLI_BASE_URL.to_owned(),
now,
expires_at,
);
Ok(Some(GrokOAuthImportEntry {
name,
email,
candidate: GrokOAuthImportCandidate::new(tokens, metadata),
}))
}
#[derive(Deserialize)]
struct OAuthCredentialsWire {
access_token: String,
refresh_token: String,
#[serde(default)]
id_token: Option<String>,
token_type: String,
expires_at: DateTime<Utc>,
email: Option<String>,
base_url: String,
client_id: String,
scope: String,
#[serde(flatten)]
extra_fields: HashMap<String, Value>,
fn credential_value<'a>(credentials: &'a Map<String, Value>, fields: &[&str]) -> Option<&'a str> {
display_value(credentials, fields).or_else(|| {
credentials
.get("token")
.and_then(Value::as_object)
.and_then(|token| display_value(token, fields))
})
}
#[derive(Default, Deserialize)]
struct OAuthAccountExtraWire {
#[serde(default)]
email: Option<String>,
#[serde(flatten)]
_extra_fields: HashMap<String, Value>,
fn display_value<'a>(object: &'a Map<String, Value>, fields: &[&str]) -> Option<&'a str> {
fields
.iter()
.find_map(|field| object.get(*field).and_then(Value::as_str))
}
fn parse_source_expiry(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(value)
.ok()
.map(|value| value.with_timezone(&Utc))
}
impl GrokOAuthImportCandidate {
backend/crates/providers/xai/src/provider.rs
+22 -1
View File
@@ -14,7 +14,8 @@ use gateway_core::engine::credential::{
};
use gateway_core::engine::provider::{
EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration,
ProviderModelCapabilities, ProviderRequest, ProviderStream, UpstreamTransport,
ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, ProviderStream,
UpstreamTransport,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt, UpstreamSendState};
use gateway_core::error::{ContinuationFailure, ProviderError, ProviderErrorKind};
@@ -114,6 +115,26 @@ impl Provider for GrokBuildProvider {
self.catalog.catalog_generation()
}
fn request_observation(&self, operation: &Operation) -> ProviderRequestObservation {
let Operation::Generate(request) = operation else {
return ProviderRequestObservation::default();
};
let reasoning_effort = request
.protocol_payload()
.body()
.get("reasoning")
.and_then(Value::as_object)
.and_then(|reasoning| reasoning.get("effort"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned);
ProviderRequestObservation {
reasoning_effort,
..ProviderRequestObservation::default()
}
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
backend/crates/providers/xai/tests/credential/import.rs
+33 -31
View File
@@ -237,49 +237,53 @@ fn oauth_account_document_should_reject_api_key_credential_field() {
}
#[test]
fn oauth_account_document_should_accept_optional_header_and_dynamic_provider_fields() {
fn oauth_account_document_should_ignore_source_wrapper_and_non_auth_metadata() {
let mut document: serde_json::Value =
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
document["type"] = serde_json::Value::String("external-account-bundle".to_owned());
document["version"] = serde_json::Value::from(1);
document["skipped_shadows"] = serde_json::Value::from(2);
document["type"] = serde_json::Value::String("sub2api-data".to_owned());
document["version"] = serde_json::Value::from(99);
document["proxies"] = serde_json::json!([{"url": "http://127.0.0.1:8080"}]);
document["source_extension"] = serde_json::json!({"opaque": true});
document["accounts"][0]["platform"] = serde_json::Value::String("source-platform".to_owned());
document["accounts"][0]["type"] = serde_json::Value::String("source-auth".to_owned());
document["accounts"][0]["concurrency"] = serde_json::Value::from(0);
document["accounts"][0]["priority"] = serde_json::Value::from(0);
document["accounts"][0]["credentials"]["provider_extension"] =
serde_json::json!({"window": "dynamic"});
document["accounts"][0]["extra"]["provider_snapshot"] = serde_json::json!({"remaining": 1});
document["accounts"][0]["extra"] = serde_json::json!({
"email": "different@example.test",
"provider_snapshot": {"remaining": 1}
});
let parsed = GrokOAuthImportDocument::parse_json(
&serde_json::to_vec(&document).expect("serialize fixture"),
)
.expect("official optional fields");
.expect("source wrapper must not constrain xAI OAuth import");
assert_eq!(parsed.into_entries().len(), 1);
}
#[test]
fn strict_oauth_account_document_should_reject_proxy_and_identity_mismatch() {
let mut proxy_document: serde_json::Value =
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
proxy_document["proxies"] = serde_json::json!([{"url": "http://127.0.0.1:8080"}]);
assert!(
GrokOAuthImportDocument::parse_json(
&serde_json::to_vec(&proxy_document).expect("serialize fixture")
)
.is_err()
);
fn minimal_oauth_credential_should_not_require_a_source_format() {
let document = serde_json::json!({
"access_token": "fixture-access-token",
"refresh_token": "fixture-refresh-token",
"expires_at": Utc::now() + Duration::hours(1)
});
let mut mismatch_document: serde_json::Value =
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
mismatch_document["accounts"][0]["extra"]["email"] =
serde_json::Value::String("different@example.test".to_owned());
let error = GrokOAuthImportDocument::parse_json(
&serde_json::to_vec(&mismatch_document).expect("serialize fixture"),
let mut entries = GrokOAuthImportDocument::parse_json(
&serde_json::to_vec(&document).expect("serialize fixture"),
)
.expect_err("identity mismatch must fail closed");
assert!(matches!(error, GrokOAuthImportError::InvalidField("email")));
.expect("minimal credential")
.into_entries();
let entry = entries.pop().expect("one entry");
assert_eq!(entry.name(), "xAI OAuth account 1");
assert_eq!(entry.email(), None);
}
#[test]
fn strict_oauth_account_document_should_reject_duplicate_account_names() {
fn oauth_account_document_should_not_require_unique_source_names() {
let mut document: serde_json::Value =
serde_json::from_slice(&oauth_account_document()).expect("fixture JSON");
let duplicate = document["accounts"][0].clone();
@@ -288,14 +292,12 @@ fn strict_oauth_account_document_should_reject_duplicate_account_names() {
.expect("accounts array")
.push(duplicate);
let error = GrokOAuthImportDocument::parse_json(
let parsed = GrokOAuthImportDocument::parse_json(
&serde_json::to_vec(&document).expect("serialize fixture"),
)
.expect_err("duplicate names must be rejected");
assert!(matches!(
error,
GrokOAuthImportError::InvalidField("account")
));
.expect("upstream identity verification owns duplicate detection");
assert_eq!(parsed.into_entries().len(), 2);
}
#[test]
backend/crates/providers/xai/tests/provider/contract.rs
+26
View File
@@ -493,6 +493,19 @@ fn operation() -> Operation {
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn operation_with_reasoning_effort(effort: &str) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("client-model")),
("input".to_owned(), json!("hello")),
("reasoning".to_owned(), json!({"effort": effort})),
]),
)
.expect("OpenAI payload");
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn compaction_operation() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
@@ -628,6 +641,19 @@ fn context_with_continuation_attempt(
context(CancellationToken::new(), Some(continuation)).with_continuation_attempt(attempt)
}
#[tokio::test]
async fn request_observation_preserves_raw_xai_reasoning_effort() {
let provider = provider(StubSelector::success(), StubInferenceTransport::success()).await;
let observation =
provider.request_observation(&operation_with_reasoning_effort("future-value"));
assert_eq!(
observation.reasoning_effort.as_deref(),
Some("future-value")
);
}
fn operation_with_state(body: serde_json::Value, state: ProviderSessionState) -> Operation {
let serde_json::Value::Object(body) = body else {
panic!("request body must be an object");
backend/migrations/0001_initial.sql
+5
View File
@@ -213,6 +213,7 @@ create table model_requests (
upstream_transport text,
http_version text,
websocket_pool text,
provider_observation_json jsonb,
attempt_count integer not null default 0,
upstream_send_state text not null default 'not_sent',
downstream_committed_at timestamptz,
@@ -316,6 +317,10 @@ create table model_requests (
constraint model_requests_websocket_pool_ck check (
websocket_pool is null or websocket_pool in ('new', 'reuse')
),
constraint model_requests_provider_observation_ck check (
provider_observation_json is null
or jsonb_typeof(provider_observation_json) = 'object'
),
constraint model_requests_cost_ck check (
cost_source in ('provider_reported', 'calculated', 'unavailable')
and (
frontend/src/views/accounts/components/AccountUsageWindow.vue
+33 -11
View File
@@ -31,6 +31,8 @@ const localUsage = computed<AccountLocalUsage | null>(() => {
return value as AccountLocalUsage
})
const mode = computed<UsageWindowMode>(() => {
if (typeof props.window?.usedPercent === 'number')
return 'quota'
if (localUsage.value)
return 'local'
if (props.window)
@@ -40,9 +42,21 @@ const mode = computed<UsageWindowMode>(() => {
const quotaValueVisible = computed(() =>
typeof props.window?.usedPercent === 'number' && props.window.usedPercent > 0,
)
const quotaLocalUsageDisplay = computed(() => {
const display = localUsage.value?.totalTokensDisplay
if (typeof display === 'string' && display.trim())
return display.trim()
const total = localUsage.value?.totalTokens
return typeof total === 'number' && total > 0 ? total.toLocaleString() : ''
})
const quotaLocalUsageVisible = computed(() =>
props.showLocalValue
&& typeof localUsage.value?.totalTokens === 'number'
&& localUsage.value.totalTokens > 0,
)
const rootClass = computed(() =>
isCompact.value ? 'min-w-0' : 'rounded-lg bg-(--cp-bg-subtle) p-2',
isCompact.value ? 'flex min-w-0 flex-col' : 'rounded-lg bg-(--cp-bg-subtle) p-2',
)
const headerClass = computed(() =>
isCompact.value
@@ -59,14 +73,14 @@ const valueClass = computed(() =>
)
const trackShapeClass = computed(() =>
isCompact.value
? 'h-1.5 w-full overflow-hidden rounded-full'
? 'h-1 w-full overflow-hidden rounded-full'
: 'h-2 overflow-hidden rounded-full',
)
const trackClass = computed(() => `${trackShapeClass.value} bg-(--cp-default-border)`)
const barStyle = computed(() => {
if (!props.window)
return undefined
return quotaWindowBarStyle(props.window, isCompact.value ? '6px' : '8px')
return quotaWindowBarStyle(props.window, isCompact.value ? '4px' : '8px')
})
const barClass = computed(() => props.window ? quotaWindowBarClass(props.window) : undefined)
const percentTextClass = computed(() =>
@@ -131,16 +145,24 @@ const requestBars = computed(() => {
<span class="min-w-0" :class="labelClass">
{{ window.labelDisplay }}
</span>
<span
v-if="quotaValueVisible"
class="shrink-0 font-mono tabular-nums"
:class="[valueClass, percentTextClass]"
>
{{ window.usedPercentDisplay }}
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
<span
v-if="quotaLocalUsageVisible"
class="text-(--cp-text-muted)"
:class="valueClass"
>
{{ quotaLocalUsageDisplay }}
</span>
<span
v-if="quotaValueVisible"
:class="[valueClass, percentTextClass]"
>
{{ window.usedPercentDisplay }}
</span>
</span>
</div>
<div
:class="[trackClass, !isCompact ? 'mt-2' : undefined]"
:class="[trackClass, !isCompact ? 'mt-2' : 'mt-auto']"
role="progressbar"
:aria-label="window.labelDisplay"
aria-valuemin="0"
@@ -176,7 +198,7 @@ const requestBars = computed(() => {
</div>
<div
class="flex items-stretch gap-px"
:class="[trackShapeClass, !isCompact ? 'mt-2' : undefined]"
:class="[trackShapeClass, !isCompact ? 'mt-2' : 'mt-auto']"
role="img"
:aria-label="requestTimelineTitle"
:title="requestTimelineTitle"
frontend/src/views/dashboard/components/AccountOverviewCard.vue
+1
View File
@@ -280,6 +280,7 @@ const valueToneClasses: Record<string, string> = {
:window="account.usageWindow"
:show-local-value="false"
variant="compact"
class="xl:h-7.25"
/>
</span>
</article>
frontend/src/views/usage/components/UsageCostCard.vue
+20 -10
View File
@@ -41,6 +41,9 @@ const props = withDefaults(
const activeView = shallowRef('cost')
const { palette } = useUsageChartPalette()
const points = computed<any[]>(() => props.cost.points)
const hasStandardCost = computed(() =>
points.value.some(point => point.standardCost != null),
)
const viewOptions = [
{ label: '费用', value: 'cost' },
@@ -101,7 +104,7 @@ function legendNames() {
return ['未缓存输入', '缓存输入', '输出']
if (activeView.value === 'cache')
return ['缓存 Token 占比', '命中请求率']
return ['预估费用', '标准费用']
return hasStandardCost.value ? ['预估费用', '标准费用'] : ['预估费用']
}
function axisFormatter() {
@@ -152,7 +155,7 @@ function chartSeries(theme: UsageChartPalette): LineSeriesOption[] {
]
}
return [
const series = [
lineSeries(
'预估费用',
chartPoints.map(point => decimalDisplayNumber(point.estimatedCost)),
@@ -160,12 +163,17 @@ function chartSeries(theme: UsageChartPalette): LineSeriesOption[] {
undefined,
true,
),
lineSeries(
'标准费用',
chartPoints.map(point => decimalDisplayNumber(point.standardCost)),
theme.textMuted,
),
]
if (hasStandardCost.value) {
series.push(
lineSeries(
'标准费用',
chartPoints.map(point => decimalDisplayNumber(point.standardCost)),
theme.textMuted,
),
)
}
return series
}
function lineSeries(
@@ -230,11 +238,13 @@ function formatTooltip(params: unknown) {
].join('<br/>')
}
return [
const lines = [
title,
`预估费用: ${formatUsd(point.estimatedCost)}`,
`标准费用: ${formatUsd(point.standardCost)}`,
].join('<br/>')
]
if (point.standardCost != null)
lines.push(`标准费用: ${formatUsd(point.standardCost)}`)
return lines.join('<br/>')
}
</script>
frontend/src/views/usage/components/UsageLatencyCell.vue
+2 -2
View File
@@ -18,8 +18,8 @@ const latencyDetails = computed(() => usageLatencyDetails(props.record))
<div
class="grid grid-cols-[auto_auto] items-center justify-end gap-x-2 gap-y-1.5 font-mono text-[12px] leading-none font-[760] tabular-nums"
>
<span class="text-[11px] text-(--cp-text-muted)">首字</span>
<span class="text-(--cp-text-secondary)">{{ latencyDetails.firstTokenDisplay }}</span>
<span class="text-[11px] text-(--cp-text-muted)">{{ latencyDetails.firstOutputLabel }}</span>
<span class="text-(--cp-text-secondary)">{{ latencyDetails.firstOutputDisplay }}</span>
<span class="text-[11px] text-(--cp-text-muted)">总耗时</span>
<span class="text-(--cp-text-primary)">{{ latencyDetails.totalDisplay }}</span>
</div>
frontend/src/views/usage/components/UsageRecordDetailModal.vue
+7 -1
View File
@@ -15,6 +15,7 @@ import {
usageBilling,
usageBillingText,
usageClientIp,
usageLatencyDetails,
usageModelDisplay,
usageReasoningEffort,
usageRecordType,
@@ -40,6 +41,7 @@ const modelDisplay = computed(() => props.record
: { primary: '—', secondary: '' })
const tokenDetails = computed(() => props.record ? usageTokenDetails(props.record) : null)
const billing = computed(() => props.record ? usageBilling(props.record) : null)
const latencyDetails = computed(() => props.record ? usageLatencyDetails(props.record) : null)
const panelClass = 'rounded-(--cp-card-radius) bg-(--cp-bg-subtle) px-4 py-3.5'
const panelTitleClass = 'm-0 text-[12px] leading-none font-[780] text-(--cp-text-secondary)'
@@ -55,7 +57,11 @@ const overviewItems = computed(() => [
{ label: '时间', value: props.record?.createdAtDisplay, mono: true },
{ label: '类型', value: props.record ? usageRecordType(props.record) : '—' },
{ label: '耗时', value: props.record?.latencyMsDisplay, mono: true },
{ label: '首字', value: props.record?.firstTokenLatencyMsDisplay, mono: true },
{
label: latencyDetails.value?.firstOutputLabel ?? '首字',
value: latencyDetails.value?.firstOutputDisplay ?? '—',
mono: true,
},
{ label: '总 Token', value: tokenDetails.value?.totalTokensDisplay, mono: true },
{ label: '请求 ID', value: props.record?.requestId, mono: true, wide: true },
{ label: '消息', value: props.record?.message, wide: true },
frontend/src/views/usage/constants.ts
+8 -3
View File
@@ -249,6 +249,7 @@ export function usageLatencyDetails(record: UsageDisplayRecord) {
const firstTokenMs = durationValue(
record.firstTokenLatencyMs ?? property(latencyDetails, 'firstTokenMs'),
)
const firstEventMs = durationValue(property(latencyDetails, 'firstEventMs'))
const totalMs = durationValue(record.latencyMs)
const firstReasoningMs = durationValue(property(latencyDetails, 'firstReasoningMs'))
const firstTextMs = durationValue(property(latencyDetails, 'firstTextMs'))
@@ -282,14 +283,16 @@ export function usageLatencyDetails(record: UsageDisplayRecord) {
},
{ label: 'WebSocket 连接', value: durationValue(property(latencyDetails, 'wsConnectMs')) },
{ label: '上游响应头', value: durationValue(property(latencyDetails, 'upstreamHeadersMs')) },
{ label: '首个上游事件', value: durationValue(property(latencyDetails, 'firstEventMs')) },
{ label: '首个上游事件', value: firstEventMs },
{ label: '上游处理', value: durationValue(property(latencyDetails, 'openaiProcessingMs')) },
]
.filter(item => item.value !== null)
.map(item => ({ ...item, value: formatDuration(item.value) }))
return {
firstTokenDisplay: formatDuration(firstTokenMs),
// SSE 可能先收到生命周期事件而没有文本或推理增量；这不是首字，须保留原始语义。
firstOutputLabel: firstTokenMs === null && firstEventMs !== null ? '首事件' : '首字',
firstOutputDisplay: formatDuration(firstTokenMs ?? firstEventMs),
totalDisplay: formatDuration(totalMs),
breakdownItems,
transportItems,
@@ -386,8 +389,10 @@ function recordMetadata(record: UsageDisplayRecord) {
}
function recordLatencyDetails(record: UsageDisplayRecord) {
const metadata = recordMetadata(record)
return asRecord('latencyDetails' in record ? record.latencyDetails : undefined)
?? recordMetadata(record)
?? asRecord(property(metadata, 'latencyDetails'))
?? metadata
}
function asRecord(value: unknown): Record<string, unknown> | undefined {