feat: revamp analytics dashboards and harden database migrations

Add dashboard and overview analytics, health monitoring, provider expense tracking, and announcement updates across the gateway and frontend.

Keep schema migrations free of historical backfills while preserving automatic backfill execution. Bound migration deadlines, run schema preparation before Compose replacement, and anonymize deleted dashboard users.

Include the current documentation cleanup and regression coverage.
This commit is contained in:
elky
2026-10-01 11:48:17 +08:00
parent 60b89cc840
commit 066ea87d72
327 changed files with 31728 additions and 20645 deletions
+2
View File
@@ -4,6 +4,7 @@ mod auth_runtime;
mod candidate_page;
mod dashboard_response;
mod direct_plan_bypass;
mod overview_total;
mod scheduler_affinity;
mod system_config;
@@ -30,6 +31,7 @@ pub(crate) use candidate_page::{
};
pub(crate) use dashboard_response::DashboardResponseCache;
pub(crate) use direct_plan_bypass::DirectPlanBypassCache;
pub(crate) use overview_total::{OverviewTotalCache, OverviewTotalRead};
pub(crate) use scheduler_affinity::{
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
};
+196
View File
@@ -0,0 +1,196 @@
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use aether_data_contracts::repository::usage::StoredUsageDashboardAnalytics;
const FRESH_FOR: Duration = Duration::from_secs(5 * 60);
const FAILURE_BACKOFF: Duration = Duration::from_secs(10);
#[derive(Debug, Default)]
pub(crate) struct OverviewTotalCache {
state: Mutex<CacheState>,
}
#[derive(Debug, Default)]
struct CacheState {
value: Option<(Instant, Arc<StoredUsageDashboardAnalytics>)>,
refreshing: bool,
retry_after: Option<Instant>,
}
pub(crate) enum OverviewTotalRead {
Pending,
Failed,
Ready {
snapshot: Arc<StoredUsageDashboardAnalytics>,
stale: bool,
},
}
/// Owns the single refresh slot even if the request that launched it disconnects.
/// Dropping a cancelled or panicking worker also releases the slot with backoff.
pub(crate) struct OverviewTotalRefresh {
cache: Arc<OverviewTotalCache>,
completed: bool,
}
impl OverviewTotalCache {
pub(crate) fn read(
self: &Arc<Self>,
now: Instant,
) -> (OverviewTotalRead, Option<OverviewTotalRefresh>) {
let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner());
let fresh = state
.value
.as_ref()
.is_some_and(|(at, _)| now.saturating_duration_since(*at) < FRESH_FOR);
let retry_allowed = state.retry_after.is_none_or(|after| now >= after);
let refresh = if !fresh && !state.refreshing && retry_allowed {
state.refreshing = true;
Some(OverviewTotalRefresh {
cache: Arc::clone(self),
completed: false,
})
} else {
None
};
let result = match &state.value {
Some((_, snapshot)) => OverviewTotalRead::Ready {
snapshot: Arc::clone(snapshot),
stale: !fresh,
},
None if state.refreshing => OverviewTotalRead::Pending,
None => OverviewTotalRead::Failed,
};
(result, refresh)
}
}
impl OverviewTotalRefresh {
pub(crate) fn finish(mut self, snapshot: Option<StoredUsageDashboardAnalytics>, now: Instant) {
let mut state = self
.cache
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
state.refreshing = false;
if let Some(snapshot) = snapshot {
state.value = Some((now, Arc::new(snapshot)));
state.retry_after = None;
} else {
state.retry_after = Some(now + FAILURE_BACKOFF);
}
self.completed = true;
}
}
impl Drop for OverviewTotalRefresh {
fn drop(&mut self) {
if !self.completed {
let mut state = self
.cache
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
state.refreshing = false;
state.retry_after = Some(Instant::now() + FAILURE_BACKOFF);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn snapshot() -> StoredUsageDashboardAnalytics {
let mut snapshot = StoredUsageDashboardAnalytics::default();
snapshot.total.generated_at = "2026-09-18T00:00:00Z".into();
snapshot.total.read_revision = "revision-1".into();
snapshot.total.summary.request_count = 42;
snapshot
}
#[test]
fn concurrent_cold_reads_claim_one_refresh() {
let cache = Arc::new(OverviewTotalCache::default());
let barrier = Arc::new(std::sync::Barrier::new(16));
let now = Instant::now();
let workers = (0..16)
.map(|_| {
let cache = Arc::clone(&cache);
let barrier = Arc::clone(&barrier);
std::thread::spawn(move || {
barrier.wait();
let (read, refresh) = cache.read(now);
assert!(matches!(read, OverviewTotalRead::Pending));
refresh
})
})
.collect::<Vec<_>>();
let mut refreshes = workers
.into_iter()
.filter_map(|worker| worker.join().unwrap())
.collect::<Vec<_>>();
assert_eq!(refreshes.len(), 1);
refreshes.pop().unwrap().finish(Some(snapshot()), now);
let (read, refresh) = cache.read(now);
assert!(matches!(
read,
OverviewTotalRead::Ready { stale: false, .. }
));
assert!(refresh.is_none());
}
#[test]
fn expiration_returns_original_snapshot_and_failed_refresh_preserves_it() {
let cache = Arc::new(OverviewTotalCache::default());
let now = Instant::now();
cache.read(now).1.unwrap().finish(Some(snapshot()), now);
assert!(cache
.read(now + FRESH_FOR - Duration::from_secs(1))
.1
.is_none());
let expired = now + FRESH_FOR;
let (read, refresh) = cache.read(expired);
let OverviewTotalRead::Ready {
snapshot: old,
stale: true,
} = read
else {
panic!("expired success must remain visible")
};
assert_eq!(old.total.generated_at, "2026-09-18T00:00:00Z");
assert_eq!(old.total.read_revision, "revision-1");
assert!(cache.read(expired).1.is_none());
refresh.unwrap().finish(None, expired);
let (read, retry) = cache.read(expired + FAILURE_BACKOFF - Duration::from_secs(1));
let OverviewTotalRead::Ready {
snapshot: retained,
stale: true,
} = read
else {
panic!("failed refresh must retain stale success")
};
assert!(Arc::ptr_eq(&old, &retained));
assert!(retry.is_none());
assert!(cache.read(expired + FAILURE_BACKOFF).1.is_some());
}
#[test]
fn cold_failure_and_worker_cancellation_back_off_before_retrying() {
let cache = Arc::new(OverviewTotalCache::default());
let now = Instant::now();
cache.read(now).1.unwrap().finish(None, now);
let (read, refresh) = cache.read(now + Duration::from_secs(9));
assert!(matches!(read, OverviewTotalRead::Failed));
assert!(refresh.is_none());
let (read, refresh) = cache.read(now + FAILURE_BACKOFF);
assert!(matches!(read, OverviewTotalRead::Pending));
drop(refresh);
let after_cancel = Instant::now();
let (read, refresh) = cache.read(after_cancel);
assert!(matches!(read, OverviewTotalRead::Failed));
assert!(refresh.is_none());
assert!(cache.read(after_cancel + FAILURE_BACKOFF).1.is_some());
}
}
@@ -7,6 +7,35 @@ pub(super) fn classify_admin_basic_family_route(
normalized_path: &str,
normalized_path_no_trailing: &str,
) -> Option<ClassifiedRoute> {
let finance_path = normalized_path_no_trailing;
if (method == http::Method::GET
&& matches!(
finance_path,
"/api/admin/billing/provider-accounts" | "/api/admin/billing/provider-expenses"
))
|| (method == http::Method::POST && finance_path == "/api/admin/billing/provider-expenses")
|| (method == http::Method::POST
&& finance_path
.strip_prefix("/api/admin/billing/provider-expenses/")
.and_then(|v| v.strip_suffix("/void"))
.is_some_and(|id| !id.is_empty() && !id.contains('/')))
{
return Some(classified(
"admin_proxy",
"billing_manage",
if finance_path.ends_with("/provider-accounts") {
"provider_accounts"
} else if method == http::Method::GET {
"provider_expenses"
} else if finance_path.ends_with("/void") {
"void_provider_expense"
} else {
"create_provider_expense"
},
"admin:billing",
false,
));
}
if method == http::Method::GET
&& matches!(
normalized_path,
@@ -6,6 +6,33 @@ pub(super) fn classify_admin_endpoints_family_route(
method: &http::Method,
normalized_path: &str,
) -> Option<ClassifiedRoute> {
if normalized_path == "/api/admin/endpoints/health/v2/publication"
&& (method == http::Method::GET || method == http::Method::PUT)
{
return Some(classified(
"admin_proxy",
"endpoints_health",
"health_v2_publication",
"admin:endpoints_health",
false,
));
}
if method == http::Method::GET
&& (matches!(
normalized_path,
"/api/admin/endpoints/health/v2/summary" | "/api/admin/endpoints/health/v2/objects"
) || normalized_path
.strip_prefix("/api/admin/endpoints/health/v2/objects/")
.is_some_and(|id| !id.is_empty() && !id.contains('/')))
{
return Some(classified(
"admin_proxy",
"endpoints_health",
"health_v2",
"admin:endpoints_health",
false,
));
}
if method == http::Method::GET && normalized_path == "/api/admin/endpoints/health/summary" {
Some(classified(
"admin_proxy",
@@ -7,6 +7,15 @@ pub(super) fn classify_admin_observability_family_route(
normalized_path: &str,
normalized_path_no_trailing: &str,
) -> Option<ClassifiedRoute> {
if let Some(kind) = classify_overview_route(method, normalized_path_no_trailing) {
return Some(classified(
"admin_proxy",
"overview_manage",
kind,
"admin:stats",
false,
));
}
if method == http::Method::POST
&& matches!(
normalized_path,
@@ -713,3 +722,32 @@ pub(super) fn classify_admin_observability_family_route(
None
}
}
fn classify_overview_route(method: &http::Method, path: &str) -> Option<&'static str> {
if method != http::Method::GET {
return None;
}
match path.strip_prefix("/api/admin/overview/")? {
"dashboard" => Some("dashboard"),
"dashboard/summary" => Some("dashboard_summary"),
"dashboard/total" => Some("dashboard_total"),
"dashboard/charts" => Some("dashboard_charts"),
"summary" => Some("summary"),
"timeseries" => Some("timeseries"),
"breakdown" => Some("breakdown"),
"users" => Some("users"),
"consumption" => Some("consumption"),
"costs" => Some("costs"),
"operations/live" => Some("operations_live"),
"operations/performance" => Some("operations_performance"),
"operations/resources" => Some("operations_resources"),
detail
if detail
.strip_prefix("users/")
.is_some_and(|id| !id.is_empty() && !id.contains('/')) =>
{
Some("user_detail")
}
_ => None,
}
}
@@ -146,6 +146,36 @@ pub(super) fn classify_public_support_route(
"public:announcements",
false,
))
} else if method == http::Method::GET
&& (matches!(
normalized_path,
"/api/users/me/health/v2/summary" | "/api/users/me/health/v2/objects"
) || normalized_path
.strip_prefix("/api/users/me/health/v2/objects/")
.is_some_and(|id| !id.is_empty() && !id.contains('/')))
{
Some(classified(
"public_support",
"health_user",
"health_v2",
"user:health",
false,
))
} else if method == http::Method::GET
&& (matches!(
normalized_path,
"/api/public/health/v2/summary" | "/api/public/health/v2/objects"
) || normalized_path
.strip_prefix("/api/public/health/v2/objects/")
.is_some_and(|id| !id.is_empty() && !id.contains('/')))
{
Some(classified(
"public_support",
"public_catalog",
"health_v2",
"public:catalog",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
@@ -273,6 +303,19 @@ pub(super) fn classify_public_support_route(
"user:monitoring",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/announcements/users/me" | "/api/announcements/users/me/"
)
{
Some(classified(
"public_support",
"announcement_user",
"list",
"user:announcements",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
@@ -203,3 +203,50 @@ fn admin_billing_plan_write_routes_buffer_request_body() {
);
}
}
#[test]
fn provider_finance_routes_require_admin_billing_and_buffer_expense_input() {
let headers = headers(&[]);
for (method, path, kind) in [
(
http::Method::GET,
"/api/admin/billing/provider-accounts",
"provider_accounts",
),
(
http::Method::GET,
"/api/admin/billing/provider-expenses",
"provider_expenses",
),
(
http::Method::POST,
"/api/admin/billing/provider-expenses",
"create_provider_expense",
),
(
http::Method::POST,
"/api/admin/billing/provider-expenses/entry-1/void",
"void_provider_expense",
),
] {
let uri: Uri = path.parse().unwrap();
let decision = classify_control_route(&method, &uri, &headers).unwrap();
assert_eq!(decision.route_family.as_deref(), Some("billing_manage"));
assert_eq!(decision.route_kind.as_deref(), Some(kind));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("admin:billing")
);
let context = GatewayPublicRequestContext::from_request_parts(
"expense-test",
&method,
&uri,
&headers,
Some(decision),
);
assert_eq!(
local_proxy_route_requires_buffered_body(&context),
kind == "create_provider_expense"
);
}
}
@@ -2,6 +2,46 @@ use http::Uri;
use super::{classify_control_route, headers};
#[test]
fn overview_routes_require_the_admin_stats_principal_and_get_method() {
for (suffix, kind) in [
("dashboard", "dashboard"),
("dashboard/summary", "dashboard_summary"),
("dashboard/total", "dashboard_total"),
("dashboard/charts", "dashboard_charts"),
("summary", "summary"),
("timeseries", "timeseries"),
("breakdown", "breakdown"),
("users", "users"),
("users/employee-1", "user_detail"),
("consumption", "consumption"),
("costs", "costs"),
("operations/live", "operations_live"),
("operations/performance", "operations_performance"),
("operations/resources", "operations_resources"),
] {
for trailing in ["", "/"] {
let uri: Uri = format!("/api/admin/overview/{suffix}{trailing}")
.parse()
.unwrap();
let decision = classify_control_route(&http::Method::GET, &uri, &headers(&[])).unwrap();
assert_eq!(decision.route_family.as_deref(), Some("overview_manage"));
assert_eq!(decision.route_kind.as_deref(), Some(kind));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("admin:stats")
);
assert!(!decision.is_execution_runtime_candidate());
let decision = classify_control_route(&http::Method::POST, &uri, &headers(&[]));
assert!(
decision.is_none_or(
|decision| decision.route_family.as_deref() != Some("overview_manage")
)
);
}
}
}
#[test]
fn classifies_admin_stats_provider_quota_usage_as_admin_proxy_route() {
let headers = headers(&[]);
@@ -261,6 +261,27 @@ fn classifies_wallet_redeem_as_public_support_route() {
);
}
#[test]
fn classifies_personal_announcements_as_authenticated_user_route() {
let headers = headers(&[]);
for path in [
"/api/announcements/users/me?limit=20&offset=0&unread_only=false",
"/api/announcements/users/me/",
] {
let uri: Uri = path.parse().expect("uri should parse");
let decision = classify_control_route(&http::Method::GET, &uri, &headers)
.expect("route should classify");
assert_eq!(decision.route_class.as_deref(), Some("public_support"));
assert_eq!(decision.route_family.as_deref(), Some("announcement_user"));
assert_eq!(decision.route_kind.as_deref(), Some("list"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("user:announcements")
);
assert!(!decision.is_execution_runtime_candidate());
}
}
#[test]
fn classifies_announcement_unread_count_as_public_support_route() {
let headers = headers(&[]);
@@ -1252,6 +1252,7 @@ impl GatewayDataState {
// exists while avoiding an unbounded read during error compensation.
let page = repository
.list_admin_wallets(&aether_data::repository::wallet::AdminWalletListQuery {
user_id: None,
status: None,
owner_type: Some("api_key".to_string()),
limit: 1,
@@ -41,6 +41,9 @@ use super::{
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
WalletLookupKey, WalletMutationOutcome,
};
use aether_data_contracts::repository::billing::{
ProviderExpenseInput, ProviderExpensePage, ProviderExpenseQuery, ProviderExpenseRecord,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
@@ -364,6 +367,26 @@ impl GatewayDataState {
}
}
pub(crate) async fn rebuild_overview_buckets(
&self,
input: &aether_data::StatsHourlyAggregationInput,
) -> Result<usize, DataLayerError> {
match &self.backends {
Some(backends) => backends.rebuild_overview_buckets(input).await,
None => Ok(0),
}
}
pub(crate) async fn drain_overview_dirty_events(
&self,
now: chrono::DateTime<chrono::Utc>,
) -> Result<u64, DataLayerError> {
match &self.backends {
Some(backends) => backends.drain_overview_dirty_events(now).await,
None => Ok(0),
}
}
pub(crate) async fn aggregate_stats_daily(
&self,
input: &aether_data::StatsDailyAggregationInput,
@@ -384,6 +407,18 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_user_announcements(
&self,
user_id: &str,
query: &aether_data::repository::announcements::UserAnnouncementListQuery,
) -> Result<aether_data::repository::announcements::StoredUserAnnouncementPage, DataLayerError>
{
match &self.announcement_reader {
Some(repository) => repository.list_user_announcements(user_id, query).await,
None => Ok(Default::default()),
}
}
pub(crate) async fn find_announcement_by_id(
&self,
announcement_id: &str,
@@ -1659,6 +1694,60 @@ impl GatewayDataState {
}
}
pub(crate) async fn query_dashboard_summary(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
) -> Result<aether_data_contracts::repository::usage::StoredDashboardSummary, DataLayerError>
{
match &self.usage_reader {
Some(repository) => repository.query_dashboard_summary(query).await,
None => Err(DataLayerError::InvalidInput(
"dashboard summary repository is unavailable".into(),
)),
}
}
pub(crate) async fn query_dashboard_analytics(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery,
) -> Result<
aether_data_contracts::repository::usage::StoredUsageDashboardAnalytics,
DataLayerError,
> {
match &self.usage_reader {
Some(repository) => repository.query_dashboard_analytics(query).await,
None => Err(DataLayerError::InvalidInput(
"usage analytics repository is unavailable".into(),
)),
}
}
pub(crate) async fn query_usage_analytics(
&self,
query: &aether_data_contracts::repository::usage::UsageAnalyticsQuery,
) -> Result<aether_data_contracts::repository::usage::StoredUsageAnalytics, DataLayerError>
{
match &self.usage_reader {
Some(repository) => repository.query_usage_analytics(query).await,
None => Err(DataLayerError::InvalidInput(
"usage analytics repository is unavailable".into(),
)),
}
}
pub(crate) async fn summarize_health_observations(
&self,
query: &aether_data_contracts::repository::usage::HealthObservationQuery,
) -> Result<aether_data_contracts::repository::usage::HealthObservationSummary, DataLayerError>
{
match &self.usage_reader {
Some(repository) => repository.summarize_health_observations(query).await,
None => Err(DataLayerError::InvalidInput(
"health observations repository is unavailable".into(),
)),
}
}
pub(crate) async fn summarize_usage_audits(
&self,
query: &aether_data_contracts::repository::usage::UsageAuditSummaryQuery,
@@ -2747,6 +2836,35 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_provider_expenses(
&self,
query: &ProviderExpenseQuery,
) -> Result<Option<ProviderExpensePage>, DataLayerError> {
match &self.billing_reader {
Some(repo) => repo.list_provider_expenses(query).await,
None => Ok(None),
}
}
pub(crate) async fn create_provider_expense(
&self,
input: &ProviderExpenseInput,
) -> Result<AdminBillingMutationOutcome<ProviderExpenseRecord>, DataLayerError> {
match &self.billing_reader {
Some(repo) => repo.create_provider_expense(input).await,
None => Ok(AdminBillingMutationOutcome::Unavailable),
}
}
pub(crate) async fn void_provider_expense(
&self,
id: &str,
operator: Option<&str>,
) -> Result<AdminBillingMutationOutcome<ProviderExpenseRecord>, DataLayerError> {
match &self.billing_reader {
Some(repo) => repo.void_provider_expense(id, operator).await,
None => Ok(AdminBillingMutationOutcome::Unavailable),
}
}
pub(crate) async fn list_billing_plans(
&self,
include_disabled: bool,
@@ -2819,6 +2937,21 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_user_plan_entitlements_with_history(
&self,
user_id: &str,
include_inactive: bool,
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.list_user_plan_entitlements_with_history(user_id, include_inactive)
.await
}
None => Ok(None),
}
}
pub(crate) async fn revoke_user_plan_entitlement(
&self,
user_id: &str,
@@ -0,0 +1,538 @@
//! Node-local, request-deduplicated activity for provider and requested-model analysis.
//!
//! RPM counts distinct requests entering upstream execution in the last 60 seconds;
//! it is never extrapolated from a shorter observation window. Concurrency follows
//! guard lifetimes, including streams, independently of that window. Expiration is
//! ordered rather than scanning request history on each lifecycle event.
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::Instant;
use chrono::{DateTime, Utc};
use serde_json::{json, Value};
const WINDOW_US: u64 = 60_000_000;
const MAX_REQUESTS: usize = 100_000;
const MAX_REQUEST_DIMENSIONS: usize = 200_000;
const MAX_LABEL_BYTES: usize = 512;
#[derive(Debug, Default)]
struct Counts {
recent: u64,
active: u64,
provider_name: Option<Arc<str>>,
}
impl Counts {
fn empty(&self) -> bool {
self.recent == 0 && self.active == 0
}
}
#[derive(Debug)]
struct ProviderRequest {
active: u64,
}
#[derive(Debug)]
struct Request {
model: Option<Arc<str>>,
active: u64,
providers: HashMap<Arc<str>, ProviderRequest>,
idle_since: Option<u64>,
cleanup_scheduled: bool,
}
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
enum Expiration {
Model(Arc<str>),
Provider(Arc<str>, Arc<str>),
Request(Arc<str>),
}
#[derive(Debug, Default)]
struct History {
through_us: u64,
requests: HashMap<Arc<str>, Request>,
providers: HashMap<Arc<str>, Counts>,
models: HashMap<Option<Arc<str>>, Counts>,
expirations: BinaryHeap<Reverse<(u64, Expiration)>>,
request_dimensions: usize,
untracked_active: u64,
incomplete_until_us: u64,
}
impl History {
fn advance(&mut self, now_us: u64) {
self.through_us = self.through_us.max(now_us);
while self
.expirations
.peek()
.is_some_and(|Reverse((expires_at, _))| *expires_at <= self.through_us)
{
let Reverse((_, expiration)) = self.expirations.pop().expect("expiration exists");
match expiration {
Expiration::Model(request_id) => {
let Some(request) = self.requests.get_mut(&request_id) else {
continue;
};
if let Some(counts) = self.models.get_mut(&request.model) {
counts.recent = counts.recent.saturating_sub(1);
if counts.empty() {
self.models.remove(&request.model);
}
}
}
Expiration::Provider(request_id, provider_id) => {
let Some(request) = self.requests.get(&request_id) else {
continue;
};
if !request.providers.contains_key(&provider_id) {
continue;
}
if let Some(counts) = self.providers.get_mut(&provider_id) {
counts.recent = counts.recent.saturating_sub(1);
if counts.empty() {
self.providers.remove(&provider_id);
}
}
}
Expiration::Request(request_id) => {
let Some(request) = self.requests.get_mut(&request_id) else {
continue;
};
request.cleanup_scheduled = false;
if let Some(idle_since) = request.idle_since {
let expires_at = idle_since.saturating_add(WINDOW_US);
if expires_at <= self.through_us {
self.request_dimensions -= request.providers.len() + 1;
self.requests.remove(&request_id);
} else {
// A retry reused the record while its first cleanup was
// pending. Keep at most one cleanup entry per request.
request.cleanup_scheduled = true;
self.expirations
.push(Reverse((expires_at, Expiration::Request(request_id))));
}
}
}
}
}
}
fn begin(
&mut self,
now_us: u64,
request_id: &str,
provider_id: &str,
provider_name: Option<&str>,
requested_model: Option<&str>,
) -> GuardIdentity {
self.advance(now_us);
let existing = self.requests.get(request_id);
let new_request = existing.is_none();
let new_provider = existing.is_none_or(|r| !r.providers.contains_key(provider_id));
let new_dimensions = usize::from(new_request) + usize::from(new_provider);
let valid_labels = !request_id.is_empty()
&& !provider_id.is_empty()
&& [
Some(request_id),
Some(provider_id),
provider_name,
requested_model,
]
.into_iter()
.flatten()
.all(|label| label.len() <= MAX_LABEL_BYTES);
if !valid_labels
|| (new_request && self.requests.len() >= MAX_REQUESTS)
|| self.request_dimensions.saturating_add(new_dimensions) > MAX_REQUEST_DIMENSIONS
{
// Telemetry must not affect admission. Explicitly mark incomplete
// coverage instead of silently returning plausible but partial counts.
self.untracked_active += 1;
self.incomplete_until_us = self.through_us.saturating_add(WINDOW_US);
return GuardIdentity::Untracked;
}
let request_id: Arc<str> = self
.requests
.get_key_value(request_id)
.map(|(key, _)| Arc::clone(key))
.unwrap_or_else(|| Arc::from(request_id));
let request = self
.requests
.entry(Arc::clone(&request_id))
.or_insert_with(|| Request {
model: requested_model
.filter(|model| !model.is_empty())
.map(Arc::from),
active: 0,
providers: HashMap::new(),
idle_since: None,
cleanup_scheduled: false,
});
let model_counts = self.models.entry(request.model.clone()).or_default();
if new_request {
model_counts.recent += 1;
self.expirations.push(Reverse((
self.through_us.saturating_add(WINDOW_US),
Expiration::Model(Arc::clone(&request_id)),
)));
}
if request.active == 0 {
model_counts.active += 1;
}
request.active += 1;
request.idle_since = None;
let provider_id: Arc<str> = request
.providers
.get_key_value(provider_id)
.map(|(key, _)| Arc::clone(key))
.unwrap_or_else(|| Arc::from(provider_id));
let provider = request
.providers
.entry(Arc::clone(&provider_id))
.or_insert(ProviderRequest { active: 0 });
let provider_counts = self.providers.entry(Arc::clone(&provider_id)).or_default();
if let Some(name) = provider_name.filter(|name| !name.is_empty()) {
provider_counts.provider_name = Some(Arc::from(name));
}
if new_provider {
provider_counts.recent += 1;
self.expirations.push(Reverse((
self.through_us.saturating_add(WINDOW_US),
Expiration::Provider(Arc::clone(&request_id), Arc::clone(&provider_id)),
)));
}
if provider.active == 0 {
provider_counts.active += 1;
}
provider.active += 1;
self.request_dimensions += new_dimensions;
GuardIdentity::Tracked {
request_id,
provider_id,
}
}
fn release(&mut self, now_us: u64, identity: GuardIdentity) {
self.advance(now_us);
let GuardIdentity::Tracked {
request_id,
provider_id,
} = identity
else {
self.untracked_active = self.untracked_active.saturating_sub(1);
return;
};
let Some(request) = self.requests.get_mut(&request_id) else {
return;
};
let Some(provider) = request.providers.get_mut(&provider_id) else {
return;
};
provider.active = provider.active.saturating_sub(1);
if provider.active == 0 {
if let Some(counts) = self.providers.get_mut(&provider_id) {
counts.active = counts.active.saturating_sub(1);
if counts.empty() {
self.providers.remove(&provider_id);
}
}
}
request.active = request.active.saturating_sub(1);
if request.active == 0 {
if let Some(counts) = self.models.get_mut(&request.model) {
counts.active = counts.active.saturating_sub(1);
if counts.empty() {
self.models.remove(&request.model);
}
}
// Retain deduplication briefly after completion as failover may begin
// after the old guard drops, including after a >60-second attempt.
request.idle_since = Some(self.through_us);
if !request.cleanup_scheduled {
request.cleanup_scheduled = true;
self.expirations.push(Reverse((
self.through_us.saturating_add(WINDOW_US),
Expiration::Request(request_id),
)));
}
}
}
fn snapshot(&mut self, now_us: u64, started_at_us: i64) -> Value {
self.advance(now_us);
let mut providers: Vec<_> = self.providers.iter().collect();
providers.sort_unstable_by(|(a, _), (b, _)| a.cmp(b));
let mut models: Vec<_> = self.models.iter().collect();
models.sort_unstable_by(|(a, _), (b, _)| a.cmp(b));
json!({
"observed_at": DateTime::from_timestamp_micros(started_at_us.saturating_add(self.through_us.min(i64::MAX as u64) as i64)),
"observed_from": DateTime::from_timestamp_micros(started_at_us),
"window_seconds": 60,
"observed_window_seconds": (self.through_us as f64 / 1_000_000.0).min(60.0),
"scope": {"kind": "node"},
"measurement": "http_and_responses_websocket_requests",
"coverage": if self.untracked_active > 0 || self.through_us < self.incomplete_until_us { "partial" } else { "complete" },
"providers": providers.into_iter().map(|(id, counts)| json!({
"provider_id": id.as_ref(),
"provider": counts.provider_name.as_deref().unwrap_or(id.as_ref()),
"requests_per_minute": counts.recent,
"current_concurrency": counts.active,
})).collect::<Vec<_>>(),
"models": models.into_iter().map(|(model, counts)| json!({
"model": model.as_deref(),
"requests_per_minute": counts.recent,
"current_concurrency": counts.active,
})).collect::<Vec<_>>(),
})
}
}
#[derive(Debug)]
pub(crate) struct ExecutionActivity {
started_at: Instant,
started_at_us: i64,
history: Mutex<History>,
}
impl Default for ExecutionActivity {
fn default() -> Self {
Self {
started_at: Instant::now(),
started_at_us: Utc::now().timestamp_micros(),
history: Mutex::new(History::default()),
}
}
}
impl ExecutionActivity {
fn elapsed_us(&self) -> u64 {
self.started_at.elapsed().as_micros().min(u64::MAX as u128) as u64
}
pub(crate) fn begin(
self: &Arc<Self>,
request_id: &str,
provider_id: &str,
provider_name: Option<&str>,
requested_model: Option<&str>,
) -> ExecutionActivityGuard {
let identity = self
.history
.lock()
.unwrap_or_else(|e| e.into_inner())
.begin(
self.elapsed_us(),
request_id,
provider_id,
provider_name,
requested_model,
);
ExecutionActivityGuard {
activity: Arc::clone(self),
identity: Some(identity),
}
}
pub(crate) fn snapshot(&self) -> Value {
self.history
.lock()
.unwrap_or_else(|e| e.into_inner())
.snapshot(self.elapsed_us(), self.started_at_us)
}
}
#[derive(Debug)]
enum GuardIdentity {
Tracked {
request_id: Arc<str>,
provider_id: Arc<str>,
},
Untracked,
}
#[derive(Debug)]
pub(crate) struct ExecutionActivityGuard {
activity: Arc<ExecutionActivity>,
identity: Option<GuardIdentity>,
}
impl Drop for ExecutionActivityGuard {
fn drop(&mut self) {
if let Some(identity) = self.identity.take() {
self.activity
.history
.lock()
.unwrap_or_else(|e| e.into_inner())
.release(self.activity.elapsed_us(), identity);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn begin(history: &mut History, at_us: u64, id: &str, provider: &str) -> GuardIdentity {
history.begin(at_us, id, provider, Some(provider), Some("requested-model"))
}
fn value(history: &mut History, at_us: u64) -> Value {
history.snapshot(at_us, 0)
}
#[test]
fn rpm_has_an_exact_rolling_window_and_never_extrapolates_startup() {
let mut history = History::default();
let a = begin(&mut history, 0, "a", "provider");
history.release(1, a);
let b = begin(&mut history, 30_000_000, "b", "provider");
history.release(30_000_001, b);
let early = value(&mut history, 30_000_001);
assert_eq!(early["providers"][0]["requests_per_minute"], 2);
assert_eq!(early["coverage"], "complete");
assert!(early["observed_window_seconds"].as_f64().unwrap() < 60.0);
assert_eq!(
value(&mut history, WINDOW_US - 1)["providers"][0]["requests_per_minute"],
2
);
assert_eq!(
value(&mut history, WINDOW_US)["providers"][0]["requests_per_minute"],
1
);
assert!(value(&mut history, 90_000_000)["providers"]
.as_array()
.unwrap()
.is_empty());
assert_eq!(
value(&mut history, 90_000_000)["observed_window_seconds"],
60.0
);
}
#[test]
fn overlapping_guards_and_sequential_retries_count_one_request() {
let mut history = History::default();
let a = begin(&mut history, 0, "request", "provider");
let b = begin(&mut history, 1, "request", "provider");
let c = begin(&mut history, 2, "other-request", "provider");
assert_eq!(
value(&mut history, 2)["providers"][0]["current_concurrency"],
2
);
assert_eq!(
value(&mut history, 2)["providers"][0]["requests_per_minute"],
2
);
history.release(3, a);
assert_eq!(
value(&mut history, 3)["models"][0]["current_concurrency"],
2
);
history.release(4, b);
history.release(5, c);
let retry = begin(&mut history, 6, "request", "provider");
let result = value(&mut history, 6);
assert_eq!(result["providers"][0]["requests_per_minute"], 2);
assert_eq!(result["models"][0]["current_concurrency"], 1);
history.release(7, retry);
}
#[test]
fn failover_counts_each_provider_but_deduplicates_the_requested_model() {
let mut history = History::default();
let first = begin(&mut history, 0, "request", "first");
let second = begin(&mut history, 1, "request", "second");
let result = value(&mut history, 2);
assert_eq!(result["providers"].as_array().unwrap().len(), 2);
assert_eq!(result["providers"][0]["requests_per_minute"], 1);
assert_eq!(result["providers"][1]["current_concurrency"], 1);
assert_eq!(result["models"][0]["requests_per_minute"], 1);
assert_eq!(result["models"][0]["current_concurrency"], 1);
history.release(3, first);
history.release(4, second);
}
#[test]
fn long_stream_retains_concurrency_and_retry_does_not_restart_model_rpm() {
let mut history = History::default();
let stream = begin(&mut history, 0, "request", "provider");
let result = value(&mut history, 2 * WINDOW_US);
assert_eq!(result["providers"][0]["requests_per_minute"], 0);
assert_eq!(result["models"][0]["current_concurrency"], 1);
history.release(2 * WINDOW_US + 1, stream);
let retry = begin(&mut history, 2 * WINDOW_US + 2, "request", "provider");
let result = value(&mut history, 2 * WINDOW_US + 2);
assert_eq!(result["providers"][0]["requests_per_minute"], 0);
assert_eq!(result["models"][0]["requests_per_minute"], 0);
assert_eq!(result["models"][0]["current_concurrency"], 1);
history.release(2 * WINDOW_US + 3, retry);
assert!(value(&mut history, 3 * WINDOW_US + 3)["models"]
.as_array()
.unwrap()
.is_empty());
assert!(history.requests.is_empty());
assert!(history.expirations.is_empty());
assert_eq!(history.request_dimensions, 0);
}
#[test]
fn cancellation_drop_releases_concurrency_but_keeps_rpm() {
let activity = Arc::new(ExecutionActivity::default());
let guard = activity.begin("request", "provider", Some("Provider name"), None);
assert_eq!(
activity.snapshot()["providers"][0]["current_concurrency"],
1
);
drop(guard);
let result = activity.snapshot();
assert_eq!(result["providers"][0]["current_concurrency"], 0);
assert_eq!(result["providers"][0]["requests_per_minute"], 1);
assert_eq!(result["providers"][0]["provider"], "Provider name");
assert!(result["models"][0]["model"].is_null());
}
#[test]
fn retry_cleanup_entries_stay_bounded_and_idle_memory_is_released() {
let mut history = History::default();
for n in 0..1_000 {
let guard = begin(&mut history, n, "request", "provider");
history.release(n, guard);
}
assert_eq!(history.expirations.len(), 3);
value(&mut history, WINDOW_US);
assert_eq!(history.expirations.len(), 1);
assert_eq!(history.requests.len(), 1);
value(&mut history, WINDOW_US + 1_000);
assert!(history.requests.is_empty());
assert!(history.providers.is_empty());
assert!(history.models.is_empty());
assert!(history.expirations.is_empty());
assert_eq!(history.request_dimensions, 0);
}
#[test]
fn sampling_limits_report_incomplete_coverage_until_unobserved_work_expires() {
let mut history = History::default();
history.request_dimensions = MAX_REQUEST_DIMENSIONS;
let untracked = begin(&mut history, 0, "request", "provider");
assert_eq!(value(&mut history, 1)["coverage"], "partial");
assert_eq!(value(&mut history, 2 * WINDOW_US)["coverage"], "partial");
history.release(2 * WINDOW_US, untracked);
assert_eq!(value(&mut history, 2 * WINDOW_US)["coverage"], "complete");
history.request_dimensions = 0;
let long_id = "x".repeat(MAX_LABEL_BYTES + 1);
let untracked = begin(&mut history, 3 * WINDOW_US, &long_id, "provider");
history.release(3 * WINDOW_US, untracked);
assert_eq!(
value(&mut history, 4 * WINDOW_US - 1)["coverage"],
"partial"
);
assert_eq!(value(&mut history, 4 * WINDOW_US)["coverage"], "complete");
}
}
@@ -181,6 +181,12 @@ async fn settle_cancelled_attempt(
usage_data.request_metadata.take(),
request_diagnostics.as_ref(),
);
usage_data.request_metadata = crate::usage::reporting::failure::with_analytics_failure(
usage_data.request_metadata.as_ref(),
"unknown",
"finalize",
"request_task_cancelled",
);
usage_data.status_code = Some(CLIENT_CANCELLED_STATUS_CODE);
usage_data.error_message = Some(error_message.to_string());
usage_data.error_category = Some("cancelled".to_string());
@@ -668,8 +668,22 @@ impl ExecutionAttemptLifecycle {
});
// 1. usage terminal
let analytics_context = if facts.provider.cancelled_by_provider() {
crate::usage::reporting::failure::with_analytics_failure(
payload.report_context.as_ref(),
"upstream",
"stream_read",
"provider_cancelled",
)
} else {
crate::usage::reporting::failure::stream_analytics_context(
payload.report_context.as_ref(),
&payload,
facts.delivery.is_aborted() && !facts.provider.is_terminal(),
)
};
let context_seed =
build_terminal_usage_context_seed(&self.plan, payload.report_context.as_ref());
build_terminal_usage_context_seed(&self.plan, analytics_context.as_ref());
let payload_seed = build_stream_terminal_usage_payload_seed(&payload);
let billing_void = settlement.billing.is_void();
let usage_runtime = Arc::clone(&state.usage_runtime);
@@ -1114,6 +1114,7 @@ fn grok_canonical_usage(usage: GrokUsageEstimate) -> StreamingCanonicalUsage {
fn grok_standardized_usage(usage: GrokUsageEstimate) -> StandardizedUsage {
let mut standardized = StandardizedUsage::new();
standardized.token_source = Some(aether_contracts::UsageTokenSource::Estimated);
standardized.input_tokens = i64::try_from(usage.input_tokens).unwrap_or(i64::MAX);
standardized.output_tokens = i64::try_from(usage.output_tokens).unwrap_or(i64::MAX);
standardized.reasoning_tokens = i64::try_from(usage.reasoning_tokens).unwrap_or(i64::MAX);
@@ -4574,6 +4575,101 @@ mod tests {
assert!(adapter.text.contains("[[1]](https://example.com/source"));
}
#[test]
fn grok_usage_reports_preserve_estimated_provenance_after_wire_roundtrip() {
use aether_usage_runtime::{
build_stream_terminal_usage_event, build_sync_terminal_usage_event,
GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEventType,
};
for (format, report_prefix) in [
("openai:chat", "openai_chat"),
("openai:responses", "openai_responses"),
] {
let mut plan = sample_plan(
serde_json::json!({
"messages": [{"role": "user", "content": "hello"}]
}),
format,
);
plan.stream = false;
plan.provider_api_format = format.to_string();
// The trusted planner binds this hint to the Grok runtime adapter.
// Exercise its transport through the same serialized report as usage.
let context = serde_json::json!({
"provider_type": "grok",
"provider_api_format": format,
"client_api_format": format,
"usage_token_source": "estimated"
});
let collected = GrokCollected {
status_code: 200,
text: "hello back".to_string(),
thinking: "short reasoning".to_string(),
..GrokCollected::default()
};
let expected = grok_usage_estimate(&plan, &collected);
let result = grok_execution_result(&plan, collected, Some(&context));
let sync_report = GatewaySyncReportRequest {
trace_id: plan.request_id.clone(),
report_kind: format!("{report_prefix}_sync_success"),
report_context: Some(context.clone()),
status_code: result.status_code,
headers: result.headers,
body_json: result.body.and_then(|body| body.json_body),
client_body_json: None,
body_base64: None,
telemetry: result.telemetry,
};
let sync_report: GatewaySyncReportRequest =
serde_json::from_slice(&serde_json::to_vec(&sync_report).unwrap()).unwrap();
let sync_event = build_sync_terminal_usage_event(
&plan,
sync_report.report_context.as_ref(),
&sync_report,
)
.unwrap();
plan.stream = true;
let stream_report = GatewayStreamReportRequest {
trace_id: plan.request_id.clone(),
report_kind: format!("{report_prefix}_stream_success"),
report_context: Some(context),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: Some(super::grok_stream_terminal_summary(&plan, expected)),
telemetry: None,
};
let stream_report: GatewayStreamReportRequest =
serde_json::from_slice(&serde_json::to_vec(&stream_report).unwrap()).unwrap();
let stream_event = build_stream_terminal_usage_event(
&plan,
stream_report.report_context.as_ref(),
&stream_report,
)
.unwrap();
// Sync honors the response's explicit total. The existing stream
// summary has no explicit total, so its fallback also adds reasoning.
let sync_total = expected.input_tokens + expected.output_tokens;
let stream_total = sync_total + expected.reasoning_tokens;
for (event, expected_total) in [(sync_event, sync_total), (stream_event, stream_total)]
{
assert_eq!(event.event_type, UsageEventType::Completed, "{format}");
assert_eq!(event.data.input_tokens, Some(expected.input_tokens));
assert_eq!(event.data.output_tokens, Some(expected.output_tokens));
assert_eq!(event.data.total_tokens, Some(expected_total));
let metadata = event.data.request_metadata.unwrap();
assert_eq!(metadata["analytics_measurement"]["source"], "estimated");
assert!(metadata.get("usage_token_source").is_none());
}
}
}
#[test]
fn openai_chat_body_includes_estimated_usage() {
let plan = sample_plan(
@@ -12,7 +12,7 @@ use std::time::{Duration, Instant};
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope};
use aether_contracts::{
ExecutionPlan, ExecutionResponseObservation, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload,
ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload, UsageTokenSource,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, UpsertRequestCandidateRecord,
@@ -445,11 +445,15 @@ fn build_sync_terminal_usage_seeds(
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
) -> (TerminalUsageContextSeed, SyncTerminalUsagePayloadSeed) {
let analytics_context =
crate::usage::reporting::failure::sync_analytics_context(report_context, payload);
let report_context_with_diagnostics =
attach_current_request_diagnostics_to_report_context(report_context);
attach_current_request_diagnostics_to_report_context(analytics_context.as_ref());
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
report_context_with_diagnostics
.as_ref()
.or(analytics_context.as_ref()),
);
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
(context_seed, payload_seed)
@@ -586,7 +590,12 @@ async fn record_stream_terminal_usage(
cancelled: bool,
) {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let context_seed = build_terminal_usage_context_seed(plan, report_context);
let analytics_context = crate::usage::reporting::failure::stream_analytics_context(
report_context,
payload,
cancelled,
);
let context_seed = build_terminal_usage_context_seed(plan, analytics_context.as_ref());
let payload_seed = build_stream_terminal_usage_payload_seed(payload);
state
.usage_runtime
@@ -976,6 +985,9 @@ async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
usage.cache_read_tokens = 0;
if usage.input_tokens <= 0 {
usage.input_tokens = estimated_input_tokens as i64;
if usage.input_tokens > 0 {
mark_kiro_stream_estimated_usage(usage, report_context, false);
}
}
return;
}
@@ -984,6 +996,10 @@ async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
usage.input_tokens = kiro_billed_input_tokens(estimated_input_tokens, cache_usage) as i64;
usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64;
usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64;
if usage.input_tokens > 0 || usage.cache_creation_tokens > 0 || usage.cache_read_tokens > 0
{
mark_kiro_stream_estimated_usage(usage, report_context, false);
}
return;
}
@@ -996,12 +1012,18 @@ async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
cache_read_input_tokens: usage.cache_read_tokens.max(0) as u64,
},
) as i64;
if usage.input_tokens > 0 {
mark_kiro_stream_estimated_usage(usage, report_context, true);
}
}
return;
}
if usage.input_tokens <= 0 {
usage.input_tokens = estimated_input_tokens as i64;
if usage.input_tokens > 0 {
mark_kiro_stream_estimated_usage(usage, report_context, true);
}
}
let Some(profile) =
@@ -1024,6 +1046,35 @@ async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
usage.input_tokens = billed_input_tokens as i64;
usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64;
usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64;
mark_kiro_stream_estimated_usage(usage, report_context, false);
}
fn mark_kiro_stream_estimated_usage(
usage: &mut StandardizedUsage,
report_context: &Value,
retains_cache: bool,
) {
let retained_source = usage.token_source.unwrap_or_else(|| {
match report_context
.get("usage_token_source")
.and_then(Value::as_str)
{
Some("estimated") => UsageTokenSource::Estimated,
Some("mixed") => UsageTokenSource::Mixed,
_ => UsageTokenSource::Reported,
}
});
let retains_reported_tokens = retained_source != UsageTokenSource::Estimated
&& (usage.output_tokens > 0
|| usage.reasoning_tokens > 0
|| usage.cache_creation_ephemeral_5m_tokens > 0
|| usage.cache_creation_ephemeral_1h_tokens > 0
|| (retains_cache && (usage.cache_creation_tokens > 0 || usage.cache_read_tokens > 0)));
usage.token_source = Some(if retains_reported_tokens {
UsageTokenSource::Mixed
} else {
UsageTokenSource::Estimated
});
}
fn append_stream_capture_bytes(
@@ -3963,7 +4014,7 @@ async fn execute_execution_runtime_stream_inner(
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let provider_in_flight_started_at = Instant::now();
let mut provider_pool_in_flight_guard =
match acquire_provider_pool_execution_guard(state, &plan).await? {
match acquire_provider_pool_execution_guard(state, &plan, report_context.as_ref()).await? {
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
record_local_runtime_candidate_skip_reason(
@@ -12293,6 +12344,10 @@ mod tests {
.expect("first usage should exist");
assert!(first_usage.cache_creation_tokens > 0);
assert_eq!(first_usage.cache_read_tokens, 0);
assert_eq!(
first_usage.token_source,
Some(aether_contracts::UsageTokenSource::Mixed)
);
let mut second_summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -12317,6 +12372,10 @@ mod tests {
assert_eq!(second_usage.cache_creation_tokens, 0);
assert!(second_usage.input_tokens < 6_000);
assert_eq!(second_usage.output_tokens, 19);
assert_eq!(
second_usage.token_source,
Some(aether_contracts::UsageTokenSource::Mixed)
);
}
#[tokio::test]
@@ -12512,6 +12571,49 @@ mod tests {
assert_eq!(usage.cache_creation_tokens, 0);
assert_eq!(usage.cache_read_tokens, 0);
assert_eq!(usage.output_tokens, 13);
assert_eq!(
usage.token_source,
Some(aether_contracts::UsageTokenSource::Mixed)
);
use aether_contracts::UsageTokenSource::{Estimated, Mixed};
for (hint, source, input, output, cache, expected) in [
(Some("estimated"), None, 0, 13, 0, Some(Estimated)),
(None, Some(Estimated), 0, 13, 0, Some(Estimated)),
(None, None, 0, 0, 200, Some(Mixed)),
(None, None, 0, 0, 0, Some(Estimated)),
(None, None, 50, 13, 0, None),
] {
let mut context = report_context.clone();
if let Some(hint) = hint {
context["usage_token_source"] = json!(hint);
}
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
token_source: source,
input_tokens: input,
output_tokens: output,
cache_read_tokens: cache,
..StandardizedUsage::new()
}),
..Default::default()
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&context),
&mut summary,
)
.await;
let usage = summary.unwrap().standardized_usage.unwrap();
assert!(usage.input_tokens > 0);
assert_eq!(usage.output_tokens, output);
assert_eq!(usage.cache_read_tokens, cache);
assert_eq!(
usage.token_source, expected,
"hint={hint:?}, source={source:?}"
);
}
}
#[tokio::test]
@@ -12751,6 +12853,10 @@ mod tests {
assert_eq!(usage.cache_creation_tokens, 175);
assert_eq!(usage.cache_read_tokens, 24_463);
assert_eq!(usage.output_tokens, 167);
assert_eq!(
usage.token_source,
Some(aether_contracts::UsageTokenSource::Mixed)
);
}
#[tokio::test]
@@ -45,6 +45,7 @@ pub(super) struct StreamFailureReport {
honor_http_failover: bool,
extra_error_fields: Map<String, Value>,
provider_body_json: Option<Value>,
analytics_failure: Option<Value>,
}
#[derive(Serialize)]
@@ -133,6 +134,7 @@ impl StreamFailureReport {
honor_http_failover: _,
mut extra_error_fields,
provider_body_json,
analytics_failure: _,
} = self;
extra_error_fields.insert("type".to_string(), Value::String(error_type));
extra_error_fields.insert("message".to_string(), Value::String(error_message));
@@ -178,6 +180,7 @@ pub(super) fn build_stream_failure_report(
honor_http_failover: false,
extra_error_fields: Map::new(),
provider_body_json: None,
analytics_failure: None,
}
}
@@ -196,6 +199,7 @@ pub(super) fn build_stream_transport_failure_report(
honor_http_failover: false,
extra_error_fields: Map::new(),
provider_body_json: None,
analytics_failure: None,
}
}
@@ -241,6 +245,10 @@ pub(super) fn build_stream_failure_from_execution_error(
honor_http_failover: error.upstream_status.is_some(),
extra_error_fields: error_object,
provider_body_json: None,
analytics_failure: crate::usage::reporting::failure::execution_error_analytics_context(
None, error,
)
.and_then(|context| context.get("analytics_failure").cloned()),
}
}
@@ -271,6 +279,7 @@ pub(super) fn build_stream_failure_from_provider_error_body(
honor_http_failover: true,
extra_error_fields: Map::new(),
provider_body_json: Some(body_json.clone()),
analytics_failure: None,
}
}
@@ -334,6 +343,7 @@ fn build_stream_failure_sync_payload(
let status_code = failure.status_code;
let upstream_status_code = failure.upstream_status_code;
let transport_error = failure.transport_error;
let analytics_failure = failure.analytics_failure.clone();
let (body, client_body) = failure.into_body_jsons();
headers.retain(|name, _| {
!name.eq_ignore_ascii_case("content-encoding")
@@ -355,6 +365,9 @@ fn build_stream_failure_sync_payload(
.or(report_context);
let report_context = report_context.map(|mut context| {
if let Some(object) = context.as_object_mut() {
if let Some(failure) = analytics_failure {
object.insert("analytics_failure".into(), failure);
}
let response_headers = serde_json::to_value(&headers).unwrap_or(Value::Null);
if upstream_status_code.is_some() {
object.insert(
@@ -499,9 +512,11 @@ async fn record_stream_sync_failure(
);
if !matches!(handling, StreamFailureHandling::HonorLocalFailover) || !retrying_next_candidate {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let analytics_context =
crate::usage::reporting::failure::sync_analytics_context(report_context, payload);
let report_context_with_diagnostics =
attach_current_request_diagnostics_and_candidate_timing_to_report_context(
report_context,
analytics_context.as_ref(),
payload
.telemetry
.as_ref()
@@ -513,7 +528,9 @@ async fn record_stream_sync_failure(
);
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
report_context_with_diagnostics
.as_ref()
.or(analytics_context.as_ref()),
);
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state
@@ -779,9 +796,13 @@ async fn handle_prefetch_transport_stream_failure(
&& matches!(analysis.decision, LocalFailoverDecision::RetryNextCandidate);
if !retrying_next_candidate {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let analytics_context = crate::usage::reporting::failure::sync_analytics_context(
payload.report_context.as_ref(),
&payload,
);
let report_context_with_diagnostics =
attach_current_request_diagnostics_and_candidate_timing_to_report_context(
payload.report_context.as_ref(),
analytics_context.as_ref(),
payload
.telemetry
.as_ref()
@@ -796,7 +817,7 @@ async fn handle_prefetch_transport_stream_failure(
plan,
report_context_with_diagnostics
.as_ref()
.or(payload.report_context.as_ref()),
.or(analytics_context.as_ref()),
);
let payload_seed = build_sync_terminal_usage_payload_seed(&payload);
state
@@ -243,7 +243,10 @@ impl SyncAttemptTerminalGuard {
record_sync_attempt_forced_terminal_state(
self.state.clone(),
self.plan.clone(),
self.report_context.clone(),
crate::usage::reporting::failure::gateway_error_analytics_context(
self.report_context.as_ref(),
error,
),
self.request_diagnostics.clone(),
self.candidate_started_unix_ms,
self.candidate_started_at,
@@ -317,6 +320,16 @@ async fn record_sync_attempt_forced_terminal_state(
let error_message = error_message.into();
let report_context =
attach_request_diagnostics_to_report_context(report_context, request_diagnostics.as_ref());
let report_context = if matches!(usage_event_type, UsageEventType::Cancelled) {
crate::usage::reporting::failure::with_analytics_failure(
report_context.as_ref(),
"unknown",
"finalize",
"request_task_cancelled",
)
} else {
report_context
};
let terminal_unix_ms = current_request_candidate_unix_ms();
let latency_ms = elapsed_ms_since(candidate_started_at);
record_local_request_candidate_status(
@@ -614,15 +627,19 @@ async fn record_sync_terminal_usage(
candidate_started_at: Instant,
candidate_first_byte_elapsed_ms: Option<u64>,
) {
let analytics_context =
crate::usage::reporting::failure::sync_analytics_context(report_context, payload);
let report_context_with_diagnostics =
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context(
report_context,
analytics_context.as_ref(),
candidate_started_at,
candidate_first_byte_elapsed_ms,
);
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
report_context_with_diagnostics
.as_ref()
.or(analytics_context.as_ref()),
);
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state
@@ -2074,37 +2091,38 @@ async fn execute_execution_runtime_sync_impl(
.unwrap_or_else(|| "-".to_string());
let candidate_started_at = Instant::now();
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let _provider_pool_in_flight_guard = match acquire_provider_pool_execution_guard(state, &plan)
.await?
{
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
record_local_runtime_candidate_skip_reason(
state,
trace_id,
"provider_key_concurrency_limit_reached",
);
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
*retry_scope = AiAttemptRetryScope::Candidate;
let _provider_pool_in_flight_guard =
match acquire_provider_pool_execution_guard(state, &plan, report_context.as_ref()).await? {
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
record_local_runtime_candidate_skip_reason(
state,
trace_id,
"provider_key_concurrency_limit_reached",
);
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
*retry_scope = AiAttemptRetryScope::Candidate;
}
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Skipped,
status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()),
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
error_message: Some(format!(
"provider key concurrency limit reached: {limit}"
)),
latency_ms: Some(0),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(candidate_started_unix_secs),
},
)
.await;
return Ok(None);
}
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Skipped,
status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()),
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
error_message: Some(format!("provider key concurrency limit reached: {limit}")),
latency_ms: Some(0),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(candidate_started_unix_secs),
},
)
.await;
return Ok(None);
}
};
};
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let usage_data = state.usage_lifecycle_data_state().as_ref().clone();
state
@@ -2804,6 +2822,9 @@ async fn execute_execution_runtime_sync_impl(
provider_response_observation.response_headers_observed_at_unix_ms,
&provider_response_observation.request_order_id,
);
if let Some(error) = result.error.as_ref() {
report_context = crate::usage::reporting::failure::execution_error_analytics_context(report_context.as_ref(), error);
}
if result.status_code >= 400 {
apply_local_execution_effect(
state,
@@ -130,6 +130,9 @@ pub(crate) async fn build_transport_error_stop_response(
None => serde_json::Map::new(),
};
request_metadata.insert("transport_error".to_string(), Value::Bool(true));
request_metadata.insert("analytics_failure".into(), json!({
"origin": "transport", "stage": "connect", "reason": "upstream_transport_error", "schema_version": 1,
}));
request_metadata.insert(
"transport_error_type".to_string(),
Value::String(error_type.to_string()),
@@ -111,6 +111,12 @@ pub(crate) fn record_failed_usage_for_deferred_response<'a>(
return;
};
let mut data = build_usage_event_data_seed(&context.plan, context.report_context.as_ref());
data.request_metadata = crate::usage::reporting::failure::with_analytics_failure(
data.request_metadata.as_ref(),
"upstream",
"response",
"candidates_exhausted",
);
data.status_code = Some(status_code);
data.error_message =
Some("all local candidates failed; returning preserved upstream error".to_string());
@@ -390,6 +396,15 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
None => Map::new(),
};
request_metadata.insert("trace_id".to_string(), Value::String(request_id.clone()));
if !request_metadata.contains_key("analytics_failure") {
request_metadata.insert(
"analytics_failure".into(),
json!({
"origin": if upstream_status_code.is_some() { "upstream" } else { "gateway" },
"stage": "routing", "reason": "candidates_exhausted", "schema_version": 1,
}),
);
}
apply_runtime_miss_usage_routing(
&mut data,
&mut request_metadata,
@@ -471,6 +486,9 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
}
let mut request_metadata = Map::new();
request_metadata.insert("analytics_failure".into(), json!({
"origin": "gateway", "stage": "routing", "reason": "execution_route_unavailable", "schema_version": 1,
}));
request_metadata.insert(
"trace_id".to_string(),
Value::String(request_id.to_string()),
@@ -16,6 +16,8 @@ mod collectors;
mod payments;
mod plans;
mod presets;
mod provider_accounts;
mod provider_expenses;
mod routes;
mod rules;
mod wallets;
@@ -207,6 +209,15 @@ pub(crate) async fn maybe_build_local_admin_billing_response(
return Ok(None);
}
if let Some(response) = provider_accounts::response(state, request_context).await? {
return Ok(Some(response));
}
if let Some(response) =
provider_expenses::response(state, request_context, request_body).await?
{
return Ok(Some(response));
}
let path = request_context.path();
let is_billing_route = (request_context.method() == http::Method::GET
&& matches!(
@@ -0,0 +1,159 @@
//! Current provider finance snapshots. This endpoint never calls upstream services.
use super::build_admin_billing_data_unavailable_response;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::{json, Value};
fn finite(value: Option<&Value>) -> Option<f64> {
value
.and_then(|v| {
v.as_f64()
.or_else(|| v.as_str().and_then(|v| v.parse::<f64>().ok()))
})
.filter(|v| v.is_finite())
}
fn text(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|v| !v.is_empty() && v.len() <= 256 && !v.chars().any(char::is_control))
}
fn timestamp(value: Option<&Value>) -> Option<String> {
let value = value?;
if let Some(raw) = value.as_str() {
if let Ok(date) = chrono::DateTime::parse_from_rfc3339(raw) {
return Some(date.to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
}
}
let secs = finite(Some(value))?;
if !(0.0..=253_402_300_799.0).contains(&secs) {
return None;
}
chrono::DateTime::from_timestamp(secs as i64, 0)
.map(|v| v.to_rfc3339_opts(chrono::SecondsFormat::Millis, true))
}
fn subscription(value: &Value) -> Value {
json!({
"group_name": text(value.get("group_name")),
"status": text(value.get("status")),
"daily_used_usd": finite(value.get("daily_used_usd")),
"daily_limit_usd": finite(value.get("daily_limit_usd")),
"weekly_used_usd": finite(value.get("weekly_used_usd")),
"weekly_limit_usd": finite(value.get("weekly_limit_usd")),
"monthly_used_usd": finite(value.get("monthly_used_usd")),
"monthly_limit_usd": finite(value.get("monthly_limit_usd")),
"expires_at": timestamp(value.get("expires_at")),
})
}
fn balance(value: &Value) -> Option<Value> {
if value.get("action_type").and_then(Value::as_str) != Some("query_balance") {
return None;
}
let status = text(value.get("status"))?;
if !matches!(status, "success" | "auth_expired" | "auth_failed") {
return None;
}
let data = value
.get("data")
.filter(|_| matches!(status, "success" | "auth_expired"));
let extra = data.and_then(|d| d.get("extra"));
let subscriptions = extra
.and_then(|e| e.get("subscriptions"))
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter(|v| v.is_object())
.take(128)
.map(subscription)
.collect::<Vec<_>>()
})
.unwrap_or_default();
Some(json!({
"status": status,
"observed_at": timestamp(value.get("executed_at")),
"currency": data.and_then(|d| text(d.get("currency"))),
"available": data.and_then(|d| finite(d.get("total_available"))),
"used": data.and_then(|d| finite(d.get("total_used"))),
"granted": data.and_then(|d| finite(d.get("total_granted"))),
"plan_name": extra.and_then(|e| text(e.get("plan_name"))),
"subscriptions": subscriptions,
}))
}
pub(super) async fn response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Option<Response<Body>>, GatewayError> {
if context.method() != http::Method::GET
|| context.path().trim_end_matches('/') != "/api/admin/billing/provider-accounts"
|| context.route_family() != Some("billing_manage")
{
return Ok(None);
}
if !state.has_provider_catalog_data_reader() {
return Ok(Some(build_admin_billing_data_unavailable_response()));
}
let mut providers = state.list_provider_catalog_providers(false).await?;
providers.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.id.cmp(&b.id)));
let keys = providers
.iter()
.map(|p| format!("provider_ops:balance:{}", p.id))
.collect::<Vec<_>>();
let (cached, unavailable) = if keys.is_empty() {
(Vec::new(), false)
} else {
match state.runtime_state().kv_get_many(&keys).await {
Ok(v) => (v, false),
Err(_) => (vec![None; keys.len()], true),
}
};
let items = providers.iter().enumerate().map(|(index, p)| {
let limit = p.monthly_quota_usd.filter(|v| v.is_finite() && *v >= 0.0);
let used = p.monthly_used_usd.filter(|v| v.is_finite() && *v >= 0.0);
let quota = if p.billing_type.as_deref() == Some("monthly_quota") || limit.is_some() {
json!({
"limit": limit, "used": used,
"remaining": limit.zip(used).map(|(l,u)| (l-u).max(0.0)),
"currency": "USD",
"period_start": p.quota_last_reset_at_unix_secs.and_then(|v| timestamp(Some(&json!(v)))),
"expires_at": p.quota_expires_at_unix_secs.and_then(|v| timestamp(Some(&json!(v)))),
})
} else { Value::Null };
let balance = cached.get(index).and_then(|v| v.as_deref())
.and_then(|v| serde_json::from_str::<Value>(v).ok()).and_then(|v| balance(&v));
json!({
"provider_id": p.id, "provider_name": p.name, "is_active": p.is_active,
"billing_type": p.billing_type, "quota": quota, "balance": balance,
})
}).collect::<Vec<_>>();
Ok(Some((
[(http::header::CACHE_CONTROL, "private, no-store")],
Json(json!({
"observed_at": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis,true),
"items": items, "balance_snapshot_unavailable": unavailable,
})),
).into_response()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn provider_accounts_only_expose_finance_allowlist_and_preserve_unknown() {
let snapshot=balance(&json!({"status":"success","action_type":"query_balance","executed_at":"2026-09-20T00:00:00Z","data":{"currency":"USD","total_available":null,"extra":{"access_token":"secret","plan_name":"Pro","subscriptions":[{"group_name":"Team","monthly_used_usd":"12.25","expires_at":1800000000,"private_token":"secret"}]}}})).unwrap();
assert!(snapshot["available"].is_null());
assert_eq!(
snapshot["subscriptions"][0]["monthly_used_usd"],
json!(12.25)
);
assert!(!snapshot.to_string().contains("secret"));
assert!(!snapshot.to_string().contains("access_token"));
let failed=balance(&json!({"status":"auth_failed","action_type":"query_balance","data":{"total_available":999}})).unwrap();
assert!(failed["available"].is_null());
}
}
@@ -0,0 +1,334 @@
use super::{
build_admin_billing_bad_request_response as bad_request,
build_admin_billing_conflict_response as conflict,
build_admin_billing_data_unavailable_response as unavailable,
build_admin_billing_not_found_response as not_found,
};
use crate::handlers::admin::{
request::{AdminAppState, AdminRequestContext},
shared::{attach_admin_audit_response, query_param_value},
};
use crate::handlers::shared::normalize_payment_currency;
use crate::GatewayError;
use aether_data_contracts::repository::billing::*;
use axum::{
body::{Body, Bytes},
http::{self, StatusCode},
response::{IntoResponse, Response},
Json,
};
use serde::Deserialize;
use serde_json::{json, Value};
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ExpenseRequest {
client_request_id: String,
provider_id: String,
kind: String,
amount: String,
currency: String,
paid_at: String,
period_start: Option<String>,
period_end: Option<String>,
note: Option<String>,
external_reference: Option<String>,
}
fn datetime(value: u64) -> String {
chrono::DateTime::from_timestamp_millis(value as i64)
.expect("valid stored timestamp")
.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
}
fn parse_date(value: &str) -> Result<u64, String> {
chrono::DateTime::parse_from_rfc3339(value)
.ok()
.and_then(|v| u64::try_from(v.timestamp_millis()).ok())
.filter(|v| *v <= 253_402_300_799_000)
.ok_or_else(|| "timestamps must be RFC3339 dates on or after 1970".into())
}
fn optional_text(value: Option<String>) -> Option<String> {
value.map(|v| v.trim().to_owned()).filter(|v| !v.is_empty())
}
fn expense_json(record: &ProviderExpenseRecord) -> Value {
let e = &record.entry;
json!({
"id": record.id, "client_request_id": e.client_request_id,
"provider_id": e.provider_id, "provider_name": e.provider_name,
"kind": e.kind, "amount": e.amount, "currency": e.currency,
"paid_at": datetime(e.paid_at_unix_ms),
"period_start": e.period_start_unix_ms.map(datetime),
"period_end": e.period_end_unix_ms.map(datetime),
"note": e.note, "external_reference": e.external_reference,
"created_by": e.created_by, "created_at": datetime(record.created_at_unix_ms),
"status": if record.voided_at_unix_ms.is_some() { "void" } else { "recorded" },
"voided_at": record.voided_at_unix_ms.map(datetime), "voided_by": record.voided_by,
})
}
fn csv_cell(value: &str) -> String {
let value = if value.trim_start().starts_with(['=', '+', '-', '@'])
|| value.starts_with(['\t', '\r', '\n'])
{
format!("'{value}")
} else {
value.to_string()
};
format!("\"{}\"", value.replace('"', "\"\""))
}
fn csv_report(items: &[ProviderExpenseRecord]) -> String {
let mut result=String::from("\u{feff}id,provider_id,provider_name,kind,amount,currency,paid_at,period_start,period_end,note,external_reference,created_by,created_at\r\n");
for r in items {
let e = &r.entry;
let fields = [
r.id.clone(),
e.provider_id.clone(),
e.provider_name.clone(),
e.kind.clone(),
e.amount.clone(),
e.currency.clone(),
datetime(e.paid_at_unix_ms),
e.period_start_unix_ms.map(datetime).unwrap_or_default(),
e.period_end_unix_ms.map(datetime).unwrap_or_default(),
e.note.clone().unwrap_or_default(),
e.external_reference.clone().unwrap_or_default(),
e.created_by.clone().unwrap_or_default(),
datetime(r.created_at_unix_ms),
];
result.push_str(
&fields
.iter()
.map(|s| csv_cell(s))
.collect::<Vec<_>>()
.join(","),
);
result.push_str("\r\n");
}
result
}
fn query(context: &AdminRequestContext<'_>, csv: bool) -> Result<ProviderExpenseQuery, String> {
let q = context.query_string();
let now = chrono::Utc::now().timestamp_millis().max(0) as u64;
let from = query_param_value(q, "from")
.map(|v| parse_date(&v))
.transpose()?
.unwrap_or(now.saturating_sub(30 * 86_400_000));
let to = query_param_value(q, "to")
.map(|v| parse_date(&v))
.transpose()?
.unwrap_or(now);
let limit = if csv {
10_001
} else {
query_param_value(q, "limit")
.map(|v| v.parse::<u32>().map_err(|_| "invalid limit".to_string()))
.transpose()?
.unwrap_or(25)
};
let offset = if csv {
0
} else {
query_param_value(q, "offset")
.map(|v| v.parse::<u64>().map_err(|_| "invalid offset".to_string()))
.transpose()?
.unwrap_or(0)
};
if !csv && limit > 200 {
return Err("limit must be at most 200".into());
}
let q = ProviderExpenseQuery {
from_unix_ms: from,
to_unix_ms: to,
limit,
offset,
};
q.validate().map_err(|e| e.to_string())?;
Ok(q)
}
pub(super) async fn response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
body: Option<&Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
let path = context.path().trim_end_matches('/');
if context.route_family() != Some("billing_manage")
|| !path.starts_with("/api/admin/billing/provider-expenses")
{
return Ok(None);
}
let operator = context
.decision()
.and_then(|d| d.admin_principal.as_ref())
.map(|p| p.user_id.clone());
if path == "/api/admin/billing/provider-expenses" && context.method() == http::Method::GET {
let csv = query_param_value(context.query_string(), "format").as_deref() == Some("csv");
let q = match query(context, csv) {
Ok(v) => v,
Err(e) => return Ok(Some(bad_request(e))),
};
let Some(page) = state
.app()
.data
.list_provider_expenses(&q)
.await
.map_err(|e| GatewayError::Internal(e.to_string()))?
else {
return Ok(Some(unavailable()));
};
if csv {
if page.total > 10_000 {
return Ok(Some(
(
StatusCode::UNPROCESSABLE_ENTITY,
Json(json!({"detail":"导出超过 10000 条,请缩小时间范围"})),
)
.into_response(),
));
}
return Ok(Some(
(
[
(http::header::CONTENT_TYPE, "text/csv; charset=utf-8"),
(
http::header::CONTENT_DISPOSITION,
"attachment; filename=provider-expenses.csv",
),
(http::header::CACHE_CONTROL, "private, no-store"),
],
csv_report(&page.items),
)
.into_response(),
));
}
return Ok(Some(
(
[(http::header::CACHE_CONTROL, "private, no-store")],
Json(json!({
"items": page.items.iter().map(expense_json).collect::<Vec<_>>(),
"total": page.total, "totals": page.totals, "providers": page.providers,
"limit": q.limit, "offset": q.offset,
"from": datetime(q.from_unix_ms), "to": datetime(q.to_unix_ms),
"time_basis": "paid_at", "source": "manual_ledger",
})),
)
.into_response(),
));
}
if path == "/api/admin/billing/provider-expenses" && context.method() == http::Method::POST {
let Some(body) = body else {
return Ok(Some(bad_request("缺少请求体")));
};
let payload = match serde_json::from_slice::<ExpenseRequest>(body) {
Ok(v) => v,
Err(_) => return Ok(Some(bad_request("输入验证失败"))),
};
let input = (|| -> Result<ProviderExpenseInput, String> {
let units = provider_expense_amount_units(&payload.amount)
.ok_or("amount must be a positive decimal string with at most 8 decimal places")?;
let input = ProviderExpenseInput {
client_request_id: uuid::Uuid::parse_str(&payload.client_request_id)
.map_err(|_| "client_request_id must be a UUID")?
.to_string(),
provider_id: payload.provider_id.trim().into(),
provider_name: "pending".into(),
kind: payload.kind,
amount: format_provider_expense_amount(units),
currency: normalize_payment_currency(&payload.currency, "currency")?,
paid_at_unix_ms: parse_date(&payload.paid_at)?,
period_start_unix_ms: payload
.period_start
.as_deref()
.map(parse_date)
.transpose()?,
period_end_unix_ms: payload.period_end.as_deref().map(parse_date).transpose()?,
note: optional_text(payload.note),
external_reference: optional_text(payload.external_reference),
created_by: operator.clone(),
};
input.validate()?;
Ok(input)
})();
let mut input = match input {
Ok(v) => v,
Err(e) => return Ok(Some(bad_request(e))),
};
let providers = state
.read_provider_catalog_providers_by_ids(&[input.provider_id.clone()])
.await?;
let Some(provider) = providers.first() else {
return Ok(Some(not_found("Provider not found")));
};
input.provider_name = provider.name.clone();
let result = state
.app()
.data
.create_provider_expense(&input)
.await
.map_err(|e| GatewayError::Internal(e.to_string()))?;
return Ok(Some(mutation_response(
result,
"admin_provider_expense_recorded",
"record_provider_expense",
)));
}
if context.method() == http::Method::POST {
if let Some(id) = path
.strip_prefix("/api/admin/billing/provider-expenses/")
.and_then(|v| v.strip_suffix("/void"))
.filter(|v| !v.is_empty() && !v.contains('/'))
{
if uuid::Uuid::parse_str(id).is_err() {
return Ok(Some(bad_request("invalid expense id")));
}
let result = state
.app()
.data
.void_provider_expense(id, operator.as_deref())
.await
.map_err(|e| GatewayError::Internal(e.to_string()))?;
return Ok(Some(mutation_response(
result,
"admin_provider_expense_voided",
"void_provider_expense",
)));
}
}
Ok(None)
}
fn mutation_response(
outcome: AdminBillingMutationOutcome<ProviderExpenseRecord>,
event: &'static str,
action: &'static str,
) -> Response<Body> {
match outcome {
AdminBillingMutationOutcome::Applied(record) => attach_admin_audit_response(
Json(json!({"item":expense_json(&record)})).into_response(),
event,
action,
"provider_expense",
&record.id,
),
AdminBillingMutationOutcome::Invalid(e) => conflict(e),
AdminBillingMutationOutcome::NotFound => not_found("Provider expense not found"),
AdminBillingMutationOutcome::Unavailable => unavailable(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn provider_expense_csv_neutralizes_formulas_and_quotes_fields() {
assert_eq!(csv_cell("=cmd()"), "\"'=cmd()\"");
assert_eq!(csv_cell(" @cmd"), "\"' @cmd\"");
assert_eq!(csv_cell("\tcmd"), "\"'\tcmd\"");
assert_eq!(csv_cell("a,\"b\"\nc"), "\"a,\"\"b\"\"\nc\"");
assert_eq!(csv_cell("12.34"), "\"12.34\"");
}
#[test]
fn provider_expense_dates_require_explicit_timezone_and_nonnegative_epoch() {
assert_eq!(
parse_date("2026-09-20T08:00:00+08:00"),
parse_date("2026-09-20T00:00:00Z")
);
assert!(parse_date("2026-09-20").is_err());
assert!(parse_date("1969-01-01T00:00:00Z").is_err());
}
}
@@ -27,11 +27,18 @@ pub(in super::super) async fn build_admin_wallet_list_response(
Ok(value) => value,
Err(detail) => return Ok(build_admin_wallets_bad_request_response(detail)),
};
let user_id = query_param_value(query, "user_id");
let status = query_param_value(query, "status");
let owner_type = parse_admin_wallets_owner_type_filter(query);
let (wallets, total) = state
.list_admin_wallets(status.as_deref(), owner_type.as_deref(), limit, offset)
.list_admin_wallets(
user_id.as_deref(),
status.as_deref(),
owner_type.as_deref(),
limit,
offset,
)
.await?;
let mut items = Vec::with_capacity(wallets.len());
for wallet in wallets {
@@ -35,11 +35,39 @@ fn build_admin_endpoint_health_bad_request_response(detail: &str) -> Response<Bo
pub(super) async fn maybe_build_local_admin_endpoints_health_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&axum::body::Bytes>,
) -> Result<Option<Response<Body>>, GatewayError> {
let Some(decision) = request_context.decision() else {
return Ok(None);
};
if decision.route_family.as_deref() == Some("endpoints_health") {
if decision.route_kind.as_deref() == Some("health_v2") {
return Ok(Some(
crate::handlers::shared::health_monitor::build_health_v2_response(
state.app(),
request_context.path(),
request_context.query_string(),
crate::handlers::shared::health_monitor::HealthAudience::Admin,
)
.await,
));
}
if decision.route_kind.as_deref() == Some("health_v2_publication") {
return Ok(Some(
crate::handlers::shared::health_monitor::build_publication_response(
state.app(),
if request_context.method() == http::Method::PUT {
Some(request_body.map_or(&[][..], |body| body.as_ref()))
} else {
None
},
)
.await,
));
}
}
if decision.route_family.as_deref() == Some("endpoints_health")
&& decision.route_kind.as_deref() == Some("health_summary")
&& request_context.path() == "/api/admin/endpoints/health/summary"
@@ -8,6 +8,7 @@ pub(crate) async fn maybe_build_local_admin_endpoints_response(
if let Some(response) = health::maybe_build_local_admin_endpoints_health_response(
&request.state(),
&request.request_context(),
request.request_body(),
)
.await?
{
@@ -1,4 +1,5 @@
mod monitoring;
mod overview;
mod routes;
mod stats;
mod usage;
@@ -22,6 +22,8 @@ pub(crate) mod test_support;
mod trace;
mod usage_helpers;
pub(super) use resilience::overview_resilience_payload;
pub(crate) async fn maybe_build_local_admin_monitoring_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -6,3 +6,23 @@ mod status;
pub(super) use history::build_admin_monitoring_resilience_circuit_history_response;
pub(super) use reset::build_admin_monitoring_reset_error_stats_response;
pub(super) use status::build_admin_monitoring_resilience_status_response;
pub(in super::super) async fn overview_resilience_payload(
state: &crate::handlers::admin::request::AdminAppState<'_>,
) -> Result<serde_json::Value, crate::GatewayError> {
let snapshot = snapshot::build_admin_monitoring_resilience_snapshot(state).await?;
let from = (snapshot.timestamp - chrono::Duration::hours(24))
.timestamp()
.max(
state
.admin_monitoring_error_stats_reset_at()
.unwrap_or_default() as i64,
);
Ok(serde_json::json!({
"scope": {"kind": "installation"},
"error_range": {"from": chrono::DateTime::from_timestamp(from, 0), "to": snapshot.timestamp},
"timestamp": snapshot.timestamp, "health_score": snapshot.health_score,
"status": snapshot.status, "error_statistics": snapshot.error_statistics,
"recent_errors": snapshot.recent_errors, "recommendations": snapshot.recommendations,
}))
}
@@ -0,0 +1,134 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::observability::analytics::{dashboard_value, parse_dashboard_query};
use axum::{
body::Body,
http::{self, StatusCode},
response::{IntoResponse, Response},
Json,
};
pub(super) async fn response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let query = match parse_dashboard_query(context.query_string()) {
Ok(query) => query,
Err(detail) => return Ok(super::error(StatusCode::BAD_REQUEST, &detail)),
};
if !state.as_ref().has_usage_data_reader() {
return Ok(super::error(
StatusCode::SERVICE_UNAVAILABLE,
"usage analytics is unavailable",
));
}
let snapshot = match tokio::time::timeout(
std::time::Duration::from_secs(15),
state.as_ref().query_dashboard_analytics(&query),
)
.await
{
Ok(result) => result?,
Err(_) => {
return Ok(super::error(
StatusCode::GATEWAY_TIMEOUT,
"dashboard query exceeded its time budget",
))
}
};
let data = dashboard_value(&query, &snapshot).map_err(GatewayError::Internal)?;
Ok((
[(http::header::CACHE_CONTROL, "private, no-store")],
Json(data),
)
.into_response())
}
pub(super) async fn total_response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
use crate::cache::OverviewTotalRead;
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
use serde_json::json;
use std::sync::Arc;
use std::time::{Duration, Instant};
let query = match parse_dashboard_query(context.query_string()) {
Ok(query) => query,
Err(detail) => return Ok(super::error(StatusCode::BAD_REQUEST, &detail)),
};
if !state.as_ref().has_usage_data_reader() {
return Ok(super::error(
StatusCode::SERVICE_UNAVAILABLE,
"usage analytics is unavailable",
));
}
let (cached, refresh) = state.as_ref().overview_total_cache.read(Instant::now());
if let Some(refresh) = refresh {
let app = state.as_ref();
let data = if app.background_data.has_usage_reader() {
Arc::clone(&app.background_data)
} else {
Arc::clone(&app.data)
};
// Lifetime boundaries do not depend on the viewer's timezone. Every
// administrator shares one refresh, including after a page reload.
tokio::spawn(async move {
let query = UsageDashboardAnalyticsQuery {
timezone: "UTC".into(),
};
let result = tokio::time::timeout(
Duration::from_secs(185),
data.query_dashboard_analytics(&query),
)
.await;
let snapshot = match result {
Ok(Ok(snapshot)) => Some(snapshot),
Ok(Err(error)) => {
tracing::warn!(%error, "dashboard lifetime refresh failed");
None
}
Err(_) => {
tracing::warn!("dashboard lifetime refresh exceeded its time budget");
None
}
};
refresh.finish(snapshot, Instant::now());
});
}
let (status, body, retry_after) = match cached {
OverviewTotalRead::Pending => {
(StatusCode::ACCEPTED, json!({"status":"pending"}), Some("3"))
}
OverviewTotalRead::Failed => (
StatusCode::SERVICE_UNAVAILABLE,
json!({"status":"failed", "detail":"cumulative dashboard totals are temporarily unavailable; retry shortly"}),
Some("10"),
),
OverviewTotalRead::Ready { snapshot, stale } => {
let mut value = dashboard_value(&query, &snapshot).map_err(GatewayError::Internal)?;
(
StatusCode::OK,
json!({
"status":"ready", "total": value["total"].take(),
"history_complete": snapshot.history_complete, "stale": stale,
}),
None,
)
}
};
let mut response = (
status,
[(http::header::CACHE_CONTROL, "private, no-store")],
Json(body),
)
.into_response();
if let Some(retry_after) = retry_after {
response.headers_mut().insert(
http::header::RETRY_AFTER,
http::HeaderValue::from_static(retry_after),
);
}
Ok(response)
}
@@ -0,0 +1,46 @@
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::observability::analytics::{dashboard_summary_value, parse_dashboard_query};
use axum::{
body::Body,
http::{header, StatusCode},
response::{IntoResponse, Response},
Json,
};
pub(super) async fn response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let query = match parse_dashboard_query(context.query_string()) {
Ok(query) => query,
Err(detail) => return Ok(super::error(StatusCode::BAD_REQUEST, &detail)),
};
if !state.as_ref().has_usage_data_reader() {
return Ok(super::error(
StatusCode::SERVICE_UNAVAILABLE,
"dashboard statistics are unavailable",
));
}
let snapshot = match tokio::time::timeout(
std::time::Duration::from_secs(5),
state.as_ref().data.query_dashboard_summary(&query),
)
.await
{
Ok(Ok(snapshot)) => snapshot,
Ok(Err(error)) => return Err(GatewayError::Internal(error.to_string())),
Err(_) => {
return Ok(super::error(
StatusCode::GATEWAY_TIMEOUT,
"dashboard statistics exceeded their time budget",
))
}
};
let mut value = dashboard_summary_value(&snapshot);
value["concurrency"] = state
.as_ref()
.today_concurrency(&query.timezone)
.map_err(GatewayError::Internal)?;
Ok(([(header::CACHE_CONTROL, "private, no-store")], Json(value)).into_response())
}
@@ -0,0 +1,144 @@
use super::error;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::observability::analytics::{envelope, metrics_value, OverviewRequest};
use aether_data_contracts::repository::usage::{UsageAnalyticsQuery, USAGE_ANALYTICS_VERSION};
use axum::{
body::Body,
http::{header, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
pub(super) async fn response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
if context
.query_string()
.is_some_and(|query| !query.is_empty())
{
return Ok(error(
StatusCode::BAD_REQUEST,
"live diagnostics do not accept historical filters",
));
}
let app = state.as_ref();
let _ = app.metric_samples().await;
let snapshot = app.metric_snapshot.read().await.clone();
let captured = snapshot.as_ref().map(|(captured, _)| *captured);
let now = chrono::Utc::now();
let observed_at = captured
.and_then(|captured| chrono::Duration::from_std(captured.elapsed()).ok())
.map(|age| now - age);
let mut unavailable = Vec::new();
let (resilience_result, recent_result) = tokio::join!(
tokio::time::timeout(
std::time::Duration::from_secs(3),
super::super::monitoring::overview_resilience_payload(state)
),
tokio::time::timeout(
std::time::Duration::from_secs(3),
recent_activity(state, now)
),
);
let resilience = match resilience_result {
Ok(Ok(value)) => Some(value),
_ => {
tracing::warn!("overview resilience snapshot unavailable");
unavailable.push("resilience");
None
}
};
let recent_activity = match recent_result {
Ok(Ok(value)) => Some(value),
_ => {
unavailable.push("recent_activity");
None
}
};
if captured.is_none() {
unavailable.push("metrics");
}
let mut response = Json(json!({
"meta": {
"schema_version": 1, "metric_version": USAGE_ANALYTICS_VERSION, "scope": {"kind": "node"},
"generated_at": now, "data_through": observed_at, "read_revision": observed_at.map(|value| value.timestamp_millis().to_string()),
"coverage": {"status": if unavailable.is_empty() {"complete"} else {"partial"}},
},
"data": {
"observed_at": observed_at, "window_seconds": null, "node_id": null,
"scope": {"kind": "node", "node_ids": []},
"metrics_text": snapshot.map(|(_, samples)| aether_runtime::metrics::render_prometheus_text(&samples)),
"resilience": resilience, "recent_activity": recent_activity,
"execution_activity": app.execution_activity.snapshot(),
"unavailable_sections": unavailable,
},
})).into_response();
response.headers_mut().insert(
header::CACHE_CONTROL,
HeaderValue::from_static("private, no-store"),
);
Ok(response)
}
async fn recent_activity(
state: &AdminAppState<'_>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<serde_json::Value, GatewayError> {
let to = now.timestamp_millis().max(60_000) as u64;
let request = OverviewRequest {
query: UsageAnalyticsQuery {
from_unix_ms: to - 60_000,
to_unix_ms: to,
timezone: "UTC".into(),
limit: 1,
..Default::default()
},
amount_basis: "billable".into(),
csv: false,
};
let snapshot = state.as_ref().query_usage_analytics(&request.query).await?;
let data = recent_activity_data(&snapshot);
Ok(envelope(&request, &snapshot, data))
}
fn recent_activity_data(
snapshot: &aether_data_contracts::repository::usage::StoredUsageAnalytics,
) -> serde_json::Value {
let mut data = metrics_value(&snapshot.summary);
data["requests_per_second"] = json!(snapshot.summary.request_count as f64 / 60.0);
data["requests_per_minute"] = json!(snapshot.summary.request_count);
data["tokens_per_minute"] = data["total_tokens"].clone();
data["window_seconds"] = json!(60);
data
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::usage::{StoredUsageAnalytics, UsageAnalyticsMetrics};
#[test]
fn recent_activity_reports_one_minute_rates_without_inventing_missing_tokens() {
let mut snapshot = StoredUsageAnalytics {
summary: UsageAnalyticsMetrics {
request_count: 120,
usage_available_count: 120,
total_tokens: 4200,
..Default::default()
},
..Default::default()
};
let value = recent_activity_data(&snapshot);
assert_eq!(value["window_seconds"], 60);
assert_eq!(value["requests_per_second"], 2.0);
assert_eq!(value["requests_per_minute"], 120);
assert_eq!(value["tokens_per_minute"], 4200);
snapshot.summary.usage_available_count = 0;
assert!(recent_activity_data(&snapshot)["tokens_per_minute"].is_null());
snapshot.summary = UsageAnalyticsMetrics::default();
assert_eq!(recent_activity_data(&snapshot)["tokens_per_minute"], 0);
}
}
@@ -0,0 +1,184 @@
mod dashboard;
mod dashboard_summary;
mod live;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::observability::analytics::{
costs_value, dashboard_charts_value, envelope, export_csv, metrics_value, page_value,
parse_dashboard_charts_query, parse_overview_query, performance_value, user_finance_value,
user_payments_value,
};
use aether_data_contracts::repository::usage::{UsageAnalyticsGranularity, UsageAnalyticsView};
use axum::{
body::Body,
http::{self, StatusCode},
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
pub(crate) async fn maybe_build_overview_response(
state: &AdminAppState<'_>,
context: &AdminRequestContext<'_>,
) -> Result<Option<Response<Body>>, GatewayError> {
if context.route_family() != Some("overview_manage") || context.method() != http::Method::GET {
return Ok(None);
}
let kind = context.route_kind().unwrap_or_default();
if kind == "dashboard_summary" {
return dashboard_summary::response(state, context).await.map(Some);
}
if kind == "dashboard_total" {
return dashboard::total_response(state, context).await.map(Some);
}
if kind == "dashboard" {
return dashboard::response(state, context).await.map(Some);
}
if matches!(kind, "operations_live" | "operations_resources") {
return live::response(state, context).await.map(Some);
}
let view = match kind {
"dashboard_charts" => UsageAnalyticsView::DashboardCharts,
"summary" => UsageAnalyticsView::Summary,
"timeseries" | "costs" => UsageAnalyticsView::Timeseries,
"operations_performance" => UsageAnalyticsView::Performance,
"users" | "user_detail" => UsageAnalyticsView::Users,
"breakdown" => UsageAnalyticsView::Breakdown,
"consumption" => UsageAnalyticsView::Consumption,
_ => return Ok(None),
};
let parsed = if kind == "dashboard_charts" {
parse_dashboard_charts_query(context.query_string())
} else {
parse_overview_query(context.query_string(), view)
};
let mut request = match parsed {
Ok(value) => value,
Err(detail) => return Ok(Some(error(StatusCode::BAD_REQUEST, &detail))),
};
if kind == "user_detail" {
let encoded = context
.path()
.trim_end_matches('/')
.rsplit('/')
.next()
.unwrap_or_default();
let Ok(id) = percent_encoding::percent_decode_str(encoded).decode_utf8() else {
return Ok(Some(error(
StatusCode::BAD_REQUEST,
"invalid user identifier",
)));
};
let id = id.as_ref();
if id.is_empty() || id.len() > 512 || id.contains('/') || id.chars().any(char::is_control) {
return Ok(Some(error(
StatusCode::BAD_REQUEST,
"invalid user identifier",
)));
}
if request
.query
.actor_user_id
.as_deref()
.is_some_and(|value| value != id)
|| request
.query
.credential_owner_id
.as_deref()
.is_some_and(|value| value != id)
{
return Ok(Some(error(
StatusCode::BAD_REQUEST,
"user filter conflicts with the requested employee",
)));
}
if request.query.attribution_kind.as_deref() == Some("employee") {
request.query.actor_user_id = Some(id.into());
} else {
request.query.credential_owner_id = Some(id.into());
}
request.query.limit = 1;
request.query.offset = 0;
}
if matches!(
view,
UsageAnalyticsView::Timeseries | UsageAnalyticsView::Performance
) {
request.query.limit = 10_000;
request.query.offset = 0;
}
if kind == "costs" {
request.query.granularity = UsageAnalyticsGranularity::Day;
}
if !state.as_ref().has_usage_data_reader() {
return Ok(Some(error(
StatusCode::SERVICE_UNAVAILABLE,
"usage analytics is unavailable",
)));
}
let snapshot = match tokio::time::timeout(
std::time::Duration::from_secs(if request.csv { 30 } else { 15 }),
state.as_ref().query_usage_analytics(&request.query),
)
.await
{
Ok(result) => result?,
Err(_) => {
return Ok(Some(error(
StatusCode::GATEWAY_TIMEOUT,
"report query exceeded its time budget; narrow the range or filters",
)))
}
};
if request.csv {
return Ok(Some(match export_csv(&request, &snapshot) {
Ok(csv) => (
[
(http::header::CONTENT_TYPE, "text/csv; charset=utf-8"),
(
http::header::CONTENT_DISPOSITION,
"attachment; filename=overview.csv",
),
(http::header::CACHE_CONTROL, "private, no-store"),
],
csv,
)
.into_response(),
Err(detail) => error(StatusCode::UNPROCESSABLE_ENTITY, &detail),
}));
}
let data = match kind {
"dashboard_charts" => dashboard_charts_value(&snapshot),
"summary" => metrics_value(&snapshot.summary),
"user_detail" => {
let Some(user) = snapshot.users.first() else {
return Ok(Some(error(StatusCode::NOT_FOUND, "employee not found")));
};
json!({
"user": { "id": user.user_id, "username": user.username, "email": user.email, "is_active": user.is_active },
"summary": metrics_value(&user.metrics),
"finance": user_finance_value(user.finance.as_ref()),
"payments": user_payments_value(snapshot.user_payments.as_ref()),
})
}
"costs" => costs_value(&request, &snapshot),
"timeseries" => {
let mut page = page_value(&request, &snapshot);
page["granularity"] = json!(request.query.granularity);
page
}
"operations_performance" => performance_value(&request, &snapshot),
_ => page_value(&request, &snapshot),
};
let mut response = Json(envelope(&request, &snapshot, data)).into_response();
response.headers_mut().insert(
http::header::CACHE_CONTROL,
http::HeaderValue::from_static("private, no-store"),
);
Ok(Some(response))
}
fn error(status: StatusCode, detail: &str) -> Response<Body> {
(status, Json(json!({"detail": detail}))).into_response()
}
@@ -1,9 +1,15 @@
use super::{monitoring, stats, usage};
use super::{monitoring, overview, stats, usage};
use crate::handlers::admin::request::{AdminRouteRequest, AdminRouteResult};
pub(crate) async fn maybe_build_local_admin_observability_response(
request: AdminRouteRequest<'_>,
) -> AdminRouteResult {
if let Some(response) =
overview::maybe_build_overview_response(&request.state(), &request.request_context())
.await?
{
return Ok(Some(response));
}
if let Some(response) =
stats::maybe_build_local_admin_stats_response(&request.state(), &request.request_context())
.await?
@@ -1,5 +1,5 @@
use super::super::super::stats::resolve_admin_usage_time_range;
use super::super::analytics::admin_usage_aggregation_by_user_json;
use super::super::summary_routes::resolve_record_time_bounds;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::query_param_value;
use crate::GatewayError;
@@ -159,12 +159,11 @@ pub(super) async fn build_admin_usage_aggregation_stats_response(
Ok(value) => value,
Err(detail) => return Ok(admin_usage_bad_request_response(detail)),
};
let time_range = match resolve_admin_usage_time_range(query) {
let time_bounds = match resolve_record_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(admin_usage_bad_request_response(detail)),
};
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
else {
let Some((created_from_unix_secs, created_until_unix_secs)) = time_bounds else {
return Ok(Json(json!([])).into_response());
};
let group_by_query = match group_by.as_str() {
@@ -33,6 +33,50 @@ use std::collections::{BTreeMap, BTreeSet};
const ADMIN_USAGE_ACTIVE_LIMIT: usize = 50;
pub(super) fn resolve_record_time_bounds(
query: Option<&str>,
) -> Result<Option<(u64, u64)>, String> {
let entries =
url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()).collect::<Vec<_>>();
let from = entries
.iter()
.filter(|(key, _)| key == "from")
.collect::<Vec<_>>();
let to = entries
.iter()
.filter(|(key, _)| key == "to")
.collect::<Vec<_>>();
if from.is_empty() && to.is_empty() {
return resolve_admin_usage_time_range(query).map(|range| range.to_unix_bounds());
}
if from.len() != 1 || to.len() != 1 {
return Err("from and to must each be provided once".into());
}
if entries
.iter()
.any(|(key, _)| matches!(key.as_ref(), "start_date" | "end_date" | "preset" | "days"))
{
return Err("precise from/to cannot be combined with date presets".into());
}
if let Some(zone) = query_param_value(query, "timezone") {
zone.parse::<chrono_tz::Tz>()
.map_err(|_| "invalid timezone".to_string())?;
}
let parse = |value: &str| -> Result<u64, String> {
let value = chrono::DateTime::parse_from_rfc3339(value)
.map_err(|_| "from/to must be RFC 3339 timestamps".to_string())?;
if value.timestamp_subsec_nanos() != 0 {
return Err("request records support second-aligned ranges".into());
}
u64::try_from(value.timestamp()).map_err(|_| "from/to must not precede Unix epoch".into())
};
let bounds = (parse(&from[0].1)?, parse(&to[0].1)?);
if bounds.0 >= bounds.1 || bounds.1 - bounds.0 > 366 * 86_400 {
return Err("from/to must define a nonempty range of at most 366 days".into());
}
Ok(Some(bounds))
}
async fn load_admin_usage_by_ids(
state: &AdminAppState<'_>,
requested_ids: &BTreeSet<String>,
@@ -70,6 +114,7 @@ fn apply_admin_usage_status_filter(query: &mut UsageAuditListQuery, status: Opti
}
"websocket" | "ws" => query.is_websocket = Some(true),
"error" | "failed" => query.error_only = true,
"success" => query.statuses = Some(vec!["completed".to_string()]),
"active" => {
query.statuses = Some(vec!["pending".to_string(), "streaming".to_string()]);
}
@@ -523,12 +568,37 @@ fn build_admin_usage_records_query(
query: Option<&str>,
limit: Option<usize>,
offset: Option<usize>,
) -> UsageAuditListQuery {
) -> Result<UsageAuditListQuery, String> {
let boolean = |key| match query_param_value(query, key).as_deref() {
None => Ok(None),
Some("true" | "1") => Ok(Some(true)),
Some("false" | "0") => Ok(Some(false)),
Some(_) => Err(format!("invalid {key}: expected true or false")),
};
let slow_threshold_ms = query_param_value(query, "slow_threshold_ms")
.map(|value| {
value
.parse::<u64>()
.ok()
.filter(|value| (1..=86_400_000).contains(value))
.ok_or_else(|| "slow_threshold_ms must be between 1 and 86400000".to_string())
})
.transpose()?;
let mut list_query = UsageAuditListQuery {
created_from_unix_secs: Some(created_from_unix_secs),
created_until_unix_secs: Some(created_until_unix_secs),
user_id: query_param_value(query, "user_id"),
provider_name: query_param_value(query, "provider"),
provider_id: query_param_value(query, "provider_id"),
api_key_id: query_param_value(query, "api_key_id"),
request_id: query_param_value(query, "request_id"),
attribution_kind: query_param_value(query, "attribution_kind"),
actor_user_id: query_param_value(query, "actor_user_id"),
slow_threshold_ms,
endpoint_kind: query_param_value(query, "endpoint_kind"),
request_type: query_param_value(query, "request_type"),
has_format_conversion: boolean("has_format_conversion")?,
is_stream: boolean("is_stream")?,
model: query_param_value(query, "model"),
api_format: query_param_value(query, "api_format"),
limit,
@@ -536,11 +606,23 @@ fn build_admin_usage_records_query(
newest_first: true,
..Default::default()
};
if list_query
.attribution_kind
.as_deref()
.is_some_and(|kind| !matches!(kind, "employee" | "standalone" | "unknown"))
{
return Err("invalid attribution_kind".into());
}
if list_query.attribution_kind.as_deref() == Some("employee")
&& list_query.actor_user_id.is_none()
{
list_query.actor_user_id = list_query.user_id.take();
}
apply_admin_usage_status_filter(
&mut list_query,
query_param_value(query, "status").as_deref(),
);
list_query
Ok(list_query)
}
fn parse_admin_usage_search_keywords(search: &str) -> Vec<String> {
@@ -653,6 +735,15 @@ fn build_admin_usage_keyword_search_query(
created_until_unix_secs: base_query.created_until_unix_secs,
user_id: base_query.user_id.clone(),
provider_name: base_query.provider_name.clone(),
provider_id: base_query.provider_id.clone(),
api_key_id: base_query.api_key_id.clone(),
request_id: base_query.request_id.clone(),
attribution_kind: base_query.attribution_kind.clone(),
actor_user_id: base_query.actor_user_id.clone(),
slow_threshold_ms: base_query.slow_threshold_ms,
endpoint_kind: base_query.endpoint_kind.clone(),
request_type: base_query.request_type.clone(),
has_format_conversion: base_query.has_format_conversion,
model: base_query.model.clone(),
api_format: base_query.api_format.clone(),
client_family: base_query.client_family.clone(),
@@ -699,13 +790,11 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
}
let query = request_context.request_query_string.as_deref();
let time_range = match resolve_admin_usage_time_range(query) {
let time_bounds = match resolve_record_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
let Some((created_from_unix_secs, created_until_unix_secs)) =
time_range.to_unix_bounds()
else {
let Some((created_from_unix_secs, created_until_unix_secs)) = time_bounds else {
return Ok(Some(build_admin_usage_summary_stats_response_from_summary(
&Default::default(),
)));
@@ -743,13 +832,11 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
sort_usage_newest_first(&mut items);
items
} else {
let time_range = match resolve_admin_usage_time_range(query) {
let time_bounds = match resolve_record_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
let Some((created_from_unix_secs, created_until_unix_secs)) =
time_range.to_unix_bounds()
else {
let Some((created_from_unix_secs, created_until_unix_secs)) = time_bounds else {
return Ok(Some(build_admin_usage_active_requests_response(
&[],
&BTreeMap::new(),
@@ -806,7 +893,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
}
let query = request_context.request_query_string.as_deref();
let time_range = match resolve_admin_usage_time_range(query) {
let time_bounds = match resolve_record_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
@@ -827,9 +914,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
let Some((created_from_unix_secs, created_until_unix_secs)) =
time_range.to_unix_bounds()
else {
let Some((created_from_unix_secs, created_until_unix_secs)) = time_bounds else {
return Ok(Some(build_admin_usage_records_response(
&[],
&BTreeMap::new(),
@@ -849,13 +934,16 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
let active_client_family_filter = client_family_filter
.as_deref()
.filter(|value| !value.trim().is_empty());
let mut base_query = build_admin_usage_records_query(
let mut base_query = match build_admin_usage_records_query(
created_from_unix_secs,
created_until_unix_secs,
query,
None,
None,
);
) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
};
base_query.client_family = active_client_family_filter.map(str::to_owned);
base_query.exclude_unknown_model_or_provider = hide_unknown_records;
let (usage, total, total_is_estimated) = if attempt_status_filter.is_some() {
@@ -1028,6 +1116,89 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
#[cfg(test)]
mod tests {
#[test]
fn precise_record_ranges_preserve_minutes_and_reject_mixed_presets() {
let range = "from=2026-09-01T23:45:00Z&to=2026-09-02T00:15:00Z&timezone=Asia%2FShanghai";
let (from, to) = super::resolve_record_time_bounds(Some(range))
.unwrap()
.unwrap();
assert_eq!(to - from, 30 * 60);
assert!(super::resolve_record_time_bounds(Some(&format!("{range}&preset=today"))).is_err());
assert!(super::resolve_record_time_bounds(Some("from=2026-09-01T00:00:00Z")).is_err());
}
#[test]
fn overview_record_drilldown_preserves_actor_and_performance_filters() {
let raw = "user_id=employee-1&attribution_kind=employee&provider_id=provider-1&api_key_id=key-1&request_id=request-1&endpoint_kind=chat&request_type=chat&is_stream=true&has_format_conversion=false&slow_threshold_ms=12000&status=success";
let query =
super::build_admin_usage_records_query(100, 200, Some(raw), None, None).unwrap();
assert_eq!(query.user_id, None);
assert_eq!(query.actor_user_id.as_deref(), Some("employee-1"));
assert_eq!(query.provider_id.as_deref(), Some("provider-1"));
assert_eq!(query.api_key_id.as_deref(), Some("key-1"));
assert_eq!(query.request_id.as_deref(), Some("request-1"));
assert_eq!(query.slow_threshold_ms, Some(12_000));
assert_eq!(query.is_stream, Some(true));
assert_eq!(query.has_format_conversion, Some(false));
assert_eq!(query.statuses, Some(vec!["completed".into()]));
let keyword = super::build_admin_usage_keyword_search_query(
&query,
vec!["example".into()],
None,
Default::default(),
false,
false,
None,
None,
);
assert_eq!(keyword.actor_user_id, query.actor_user_id);
assert_eq!(keyword.slow_threshold_ms, query.slow_threshold_ms);
assert_eq!(keyword.has_format_conversion, query.has_format_conversion);
for invalid in [
"is_stream=maybe",
"slow_threshold_ms=0",
"attribution_kind=owner",
] {
assert!(
super::build_admin_usage_records_query(100, 200, Some(invalid), None, None)
.is_err()
);
}
}
#[test]
fn overview_record_drilldown_preserves_standalone_key_ownership() {
let raw = "user_id=owner-1&attribution_kind=standalone&api_key_id=standalone-key";
let query =
super::build_admin_usage_records_query(100, 200, Some(raw), None, None).unwrap();
assert_eq!(query.attribution_kind.as_deref(), Some("standalone"));
assert_eq!(query.user_id.as_deref(), Some("owner-1"));
assert_eq!(query.actor_user_id, None);
assert_eq!(query.api_key_id.as_deref(), Some("standalone-key"));
let keyword = super::build_admin_usage_keyword_search_query(
&query,
vec!["example".into()],
None,
Default::default(),
false,
false,
None,
None,
);
assert_eq!(keyword.attribution_kind, query.attribution_kind);
assert_eq!(keyword.user_id, query.user_id);
assert_eq!(keyword.actor_user_id, query.actor_user_id);
assert_eq!(keyword.api_key_id, query.api_key_id);
for retired_kind in ["service", "shared"] {
let raw = format!("attribution_kind={retired_kind}");
assert!(
super::build_admin_usage_records_query(100, 200, Some(&raw), None, None).is_err(),
"{retired_kind}"
);
}
}
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
@@ -1169,7 +1340,7 @@ mod tests {
for status in ["websocket", "ws", "WS"] {
let raw_query = format!("status={status}");
let list_query =
build_admin_usage_records_query(100, 200, Some(&raw_query), None, None);
build_admin_usage_records_query(100, 200, Some(&raw_query), None, None).unwrap();
assert_eq!(list_query.is_websocket, Some(true));
assert_eq!(list_query.is_stream, None);
@@ -1190,7 +1361,7 @@ mod tests {
for (status, expected_stream) in [("stream", true), ("standard", false)] {
let raw_query = format!("status={status}");
let list_query =
build_admin_usage_records_query(100, 200, Some(&raw_query), None, None);
build_admin_usage_records_query(100, 200, Some(&raw_query), None, None).unwrap();
assert_eq!(list_query.is_stream, Some(expected_stream));
assert_eq!(list_query.is_websocket, Some(false));
@@ -121,6 +121,7 @@ impl<'a> AdminAppState<'a> {
pub(crate) async fn list_admin_wallets(
&self,
user_id: Option<&str>,
status: Option<&str>,
owner_type: Option<&str>,
limit: usize,
@@ -133,7 +134,7 @@ impl<'a> AdminAppState<'a> {
GatewayError,
> {
self.app
.list_admin_wallets(status, owner_type, limit, offset)
.list_admin_wallets(user_id, status, owner_type, limit, offset)
.await
}
@@ -1,7 +1,9 @@
use super::{build_admin_users_bad_request_response, build_admin_users_data_unavailable_response};
use crate::handlers::admin::billing::admin_payment_gateway_response_projection;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{attach_admin_audit_response, unix_secs_to_rfc3339};
use crate::handlers::admin::shared::{
attach_admin_audit_response, query_param_value, unix_secs_to_rfc3339,
};
use crate::handlers::shared::unix_ms_to_rfc3339;
use crate::GatewayError;
use aether_data_contracts::repository::billing::{BillingPlanRecord, UserPlanEntitlementRecord};
@@ -177,8 +179,13 @@ fn entitlement_payload(
async fn load_admin_user_entitlements_payload(
state: &AdminAppState<'_>,
user_id: &str,
include_inactive: bool,
) -> Result<Option<serde_json::Value>, GatewayError> {
let entitlements = match state.app().list_user_plan_entitlements(user_id).await? {
let entitlements = match state
.app()
.list_user_plan_entitlements_with_history(user_id, include_inactive)
.await?
{
Some(value) => value,
None => return Ok(None),
};
@@ -214,7 +221,17 @@ pub(in super::super) async fn build_admin_list_user_billing_entitlements_respons
)
.into_response());
}
match load_admin_user_entitlements_payload(state, &user_id).await? {
let include_inactive =
match query_param_value(request_context.query_string(), "include_inactive").as_deref() {
None | Some("false" | "0") => false,
Some("true" | "1") => true,
_ => {
return Ok(build_admin_users_bad_request_response(
"include_inactive 必须为布尔值",
))
}
};
match load_admin_user_entitlements_payload(state, &user_id, include_inactive).await? {
Some(payload) => Ok(Json(payload).into_response()),
None => Ok(build_admin_users_data_unavailable_response()),
}
@@ -256,7 +273,7 @@ pub(in super::super) async fn build_admin_revoke_user_billing_entitlement_respon
return Ok(build_admin_users_data_unavailable_response());
}
}
let entitlements = match load_admin_user_entitlements_payload(state, &user_id).await? {
let entitlements = match load_admin_user_entitlements_payload(state, &user_id, false).await? {
Some(value) => value,
None => return Ok(build_admin_users_data_unavailable_response()),
};
@@ -401,7 +418,7 @@ pub(in super::super) async fn build_admin_grant_user_billing_plan_response(
return Ok(build_admin_users_data_unavailable_response());
}
};
let entitlements = match load_admin_user_entitlements_payload(state, &user_id).await? {
let entitlements = match load_admin_user_entitlements_payload(state, &user_id, false).await? {
Some(value) => value,
None => return Ok(build_admin_users_data_unavailable_response()),
};
+15 -2
View File
@@ -1996,9 +1996,22 @@ async fn proxy_request_inner(
request_permit = aether_runtime::AdmissionPermit::combine(
request_permit.into_iter().chain(plan_usage_permit),
);
if let Some(request_permit) = request_permit.as_ref() {
// The affinity-forwarding node already returned above. Observe only local
// AI execution, retaining the guard in both the body and detached execution.
let activity_permit = control_decision
.is_some_and(|decision| {
decision.route_class.as_deref() == Some("ai_public")
&& decision.execution_runtime_candidate
})
.then(|| state.request_activity.begin().into_admission_permit());
if let Some(activity) = activity_permit.as_ref() {
crate::request_lifecycle::track_request_activity(activity.clone());
}
if let Some(background_permit) = aether_runtime::AdmissionPermit::combine(
request_permit.clone().into_iter().chain(activity_permit),
) {
parts.extensions.insert(
crate::executor::candidate_loop::BackgroundAdmissionPermit::new(request_permit.clone()),
crate::executor::candidate_loop::BackgroundAdmissionPermit::new(background_permit),
);
}
@@ -376,6 +376,12 @@ async fn handle_live_http(
state,
&attempt.plan,
request_context.trace_id.as_str(),
Some(
attempt
.report_context
.as_ref()
.unwrap_or(&serde_json::Value::Null),
),
)
.await
{
@@ -1140,8 +1140,13 @@ async fn acquire_live_relay_admission(
});
}
}
match ResponsesWebSocketTurnAdmission::acquire(state, &attempt.plan, context.trace_id.as_str())
.await
match ResponsesWebSocketTurnAdmission::acquire(
state,
&attempt.plan,
context.trace_id.as_str(),
None,
)
.await
{
Ok(capacity) => Ok(LiveRelayAdmission { capacity, audit }),
Err(error) => Err(LiveRelayAdmissionFailure {
@@ -136,6 +136,7 @@ pub(super) async fn prepare_realtime_websocket(
state,
&candidate.admission_plan,
context.trace_id.as_str(),
None,
)
.await
{
@@ -11,7 +11,8 @@ use aether_contracts::ExecutionPlan;
use crate::execution_runtime::acquire_upstream_execution_gate;
use crate::provider_pool_demand::{
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard,
acquire_provider_pool_execution_guard, acquire_provider_pool_execution_guard_unobserved,
ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard,
};
use crate::upstream_admission::UpstreamTargetAdmissionPermit;
use crate::{AppState, GatewayError};
@@ -28,6 +29,7 @@ impl ResponsesWebSocketTurnAdmission {
state: &AppState,
plan: &ExecutionPlan,
trace_id: &str,
observation_context: Option<&serde_json::Value>,
) -> Result<Self, GatewayError> {
let upstream_execution = acquire_upstream_execution_gate(state, trace_id).await?;
let upstream_target = match state
@@ -41,7 +43,13 @@ impl ResponsesWebSocketTurnAdmission {
return Err(error);
}
};
let provider_pool = match acquire_provider_pool_execution_guard(state, plan).await? {
let provider_admission = match observation_context {
Some(context) => {
acquire_provider_pool_execution_guard(state, plan, Some(context)).await?
}
None => acquire_provider_pool_execution_guard_unobserved(state, plan).await?,
};
let provider_pool = match provider_admission {
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
drop(upstream_target);
@@ -32,6 +32,8 @@ pub(super) struct ResponsesWebSocketTurnControl {
pub(super) decision: GatewayControlDecision,
pub(super) auth_snapshot: Option<GatewayAuthApiKeySnapshot>,
pub(super) rpm_bypassed: bool,
// Shared across transparent retries and owned by LogicalTurn, not the socket.
pub(super) activity: std::sync::Arc<crate::request_activity::RequestActivityGuard>,
}
pub(super) async fn resolve_responses_websocket_turn_control(
@@ -124,6 +126,7 @@ pub(super) async fn resolve_responses_websocket_turn_control(
decision,
auth_snapshot,
rpm_bypassed,
activity: std::sync::Arc::new(state.request_activity.begin()),
})
}
@@ -606,6 +606,7 @@ pub(super) async fn begin_unowned_responses_websocket_turn(
state,
&plan,
plan.request_id.as_str(),
Some(report_context.as_ref().unwrap_or(&Value::Null)),
)
.await
{
@@ -260,6 +260,52 @@ mod tests {
)
}
#[test]
fn concurrency_activity_survives_retry_and_ends_with_the_logical_turn() {
let activity = std::sync::Arc::new(crate::request_activity::RequestActivity::default());
let first = logical().with_turn_control(super::ResponsesWebSocketTurnControl {
decision: crate::control::GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".into()),
None,
None,
None,
),
auth_snapshot: None,
rpm_bypassed: false,
activity: std::sync::Arc::new(activity.begin()),
});
let mut state = ResponsesTurnState::Idle;
state.begin(first, FakeAttempt(1));
assert_eq!(activity.active(), 1);
assert_eq!(state.detach_attempt(), Some(FakeAttempt(1)));
assert_eq!(
activity.active(),
1,
"transparent retry retains one logical request"
);
state.resume(FakeAttempt(2)).unwrap();
assert_eq!(activity.active(), 1);
assert_eq!(state.end(), Some(FakeAttempt(2)));
assert_eq!(activity.active(), 0);
let logical = logical().with_turn_control(super::ResponsesWebSocketTurnControl {
decision: crate::control::GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".into()),
None,
None,
None,
),
auth_snapshot: None,
rpm_bypassed: false,
activity: std::sync::Arc::new(activity.begin()),
});
state.begin(logical, FakeAttempt(3));
drop(state);
assert_eq!(activity.active(), 0, "disconnect drops the logical request");
}
/// 透明重试失败之后:旧 attempt 已经被 detach 并结算过,logical turn 仍停在
/// `Replanning`。此时 `end()` 不能再交出 attempt,否则同一个 attempt 会被
/// 结算两次(两条 usage terminal、两次 pool lease 释放)。
@@ -166,6 +166,23 @@ async fn build_local_public_support_response(
return None;
}
if decision.route_family.as_deref() == Some("health_user") {
if let Err(response) =
resolve_authenticated_local_user(state, request_context, headers).await
{
return Some(response);
}
return Some(
crate::handlers::shared::health_monitor::build_health_v2_response(
state,
&request_context.request_path,
request_context.request_query_string.as_deref(),
crate::handlers::shared::health_monitor::HealthAudience::Authenticated,
)
.await,
);
}
if decision.route_family.as_deref() == Some("auth") {
return maybe_build_local_auth_response(
state,
@@ -306,6 +323,17 @@ async fn build_local_public_support_response(
}
if decision.route_family.as_deref() == Some("public_catalog") {
if decision.route_kind.as_deref() == Some("health_v2") {
return Some(
crate::handlers::shared::health_monitor::build_health_v2_response(
state,
&request_context.request_path,
request_context.request_query_string.as_deref(),
crate::handlers::shared::health_monitor::HealthAudience::Public,
)
.await,
);
}
if decision.route_kind.as_deref() == Some("site_info")
&& request_context.request_path == "/api/public/site-info"
{
@@ -1,3 +1,4 @@
use aether_data::repository::announcements::UserAnnouncementListQuery;
use axum::{
body::Body,
http,
@@ -33,6 +34,40 @@ fn parse_announcement_read_status_request(
}
}
fn parse_user_announcements_query(
raw: Option<&str>,
now_unix_secs: u64,
) -> Result<UserAnnouncementListQuery, String> {
let mut query = UserAnnouncementListQuery {
unread_only: false,
offset: 0,
limit: 20,
now_unix_secs,
};
let mut seen = std::collections::BTreeSet::new();
for (key, value) in url::form_urlencoded::parse(raw.unwrap_or_default().as_bytes()) {
if !seen.insert(key.clone()) {
return Err(format!("duplicate announcement query parameter: {key}"));
}
match key.as_ref() {
"limit" => {
query.limit = value.parse().map_err(|_| "invalid announcement limit")?;
}
"offset" => {
query.offset = value.parse().map_err(|_| "invalid announcement offset")?;
}
"unread_only" => {
query.unread_only = value
.parse()
.map_err(|_| "unread_only must be true or false")?;
}
_ => return Err(format!("unsupported announcement query parameter: {key}")),
}
}
query.validate().map_err(|err| err.to_string())?;
Ok(query)
}
pub(crate) async fn maybe_build_local_announcement_user_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
@@ -54,6 +89,48 @@ pub(crate) async fn maybe_build_local_announcement_user_response(
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
match decision.route_kind.as_deref() {
Some("list")
if request_context.request_method == http::Method::GET
&& matches!(
request_context.request_path.as_str(),
"/api/announcements/users/me" | "/api/announcements/users/me/"
) =>
{
let query = match parse_user_announcements_query(
request_context.request_query_string.as_deref(),
now_unix_secs,
) {
Ok(query) => query,
Err(detail) => return Some(announcements_bad_request_response(detail)),
};
let page = match state.list_user_announcements(&auth.user.id, &query).await {
Ok(page) => page,
Err(err) => {
return Some(announcements_internal_error_response(
announcements_internal_detail(err),
))
}
};
let items = page
.items
.iter()
.map(|item| {
let mut value = build_public_announcement_payload(&item.announcement);
value["is_read"] = json!(item.is_read);
value
})
.collect::<Vec<_>>();
Some(
Json(json!({
"items": items,
"total": page.total,
"unread_count": page.unread_count,
"limit": query.limit,
"offset": query.offset,
}))
.into_response(),
)
}
Some("unread_count")
if request_context.request_method == http::Method::GET
&& matches!(
@@ -1209,6 +1209,7 @@ pub(super) async fn handle_users_me_usage_get(
limit: None,
offset: None,
newest_first: true,
..Default::default()
};
total_record_count = match state
.count_usage_audits_by_keyword_search(&keyword_query)
@@ -1259,6 +1260,7 @@ pub(super) async fn handle_users_me_usage_get(
limit: None,
offset: None,
newest_first: true,
..Default::default()
})
.await
{
@@ -1289,6 +1291,7 @@ pub(super) async fn handle_users_me_usage_get(
limit: Some(limit),
offset: Some(offset),
newest_first: true,
..Default::default()
})
.await
{
@@ -1434,6 +1437,7 @@ pub(super) async fn handle_users_me_usage_active_get(
limit: Some(50),
offset: None,
newest_first: true,
..Default::default()
})
.await
{
@@ -0,0 +1,511 @@
use super::policy::{overall_status, HealthPolicy, HealthRatio, HealthStatus};
use super::publication::{HealthPublication, PUBLICATION_KEY};
use crate::handlers::shared::unix_ms_to_rfc3339;
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::global_models::AdminGlobalModelListQuery;
use aether_data_contracts::repository::usage::{
HealthObservationMetrics, HealthObservationObjectKind, HealthObservationQuery,
};
use axum::{body::Body, http::StatusCode, response::IntoResponse, response::Response, Json};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use chrono::Utc;
use serde::Serialize;
use serde_json::json;
use std::collections::BTreeMap;
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum HealthAudience {
Admin,
Authenticated,
Public,
}
#[derive(Clone, Debug)]
pub(super) struct HealthRequest {
pub kind: HealthObservationObjectKind,
pub from_unix_ms: u64,
pub to_unix_ms: u64,
pub limit: usize,
pub offset: usize,
}
impl HealthRequest {
pub fn parse(query: Option<&str>, public: bool, now_ms: u64) -> Result<Self, &'static str> {
let mut params = BTreeMap::new();
for (key, value) in url::form_urlencoded::parse(query.unwrap_or_default().as_bytes()) {
if !matches!(key.as_ref(), "kind" | "window" | "limit" | "offset") {
return Err("Unsupported health query parameter");
}
if params
.insert(key.into_owned(), value.into_owned())
.is_some()
{
return Err("Duplicate health query parameter");
}
}
let kind = match params
.get("kind")
.map(String::as_str)
.unwrap_or("api_format")
{
"api_format" => HealthObservationObjectKind::ApiFormat,
"model" => HealthObservationObjectKind::Model,
"provider" if !public => HealthObservationObjectKind::Provider,
_ => return Err("Unsupported health object kind"),
};
let hours = match params.get("window").map(String::as_str).unwrap_or("6h") {
"1h" => 1,
"6h" => 6,
"24h" => 24,
"72h" => 72,
_ => return Err("Health window must be 1h, 6h, 24h or 72h"),
};
let parse_size = |name: &str, default: usize| -> Result<usize, &'static str> {
params
.get(name)
.map(|value| value.parse().map_err(|_| "Invalid pagination value"))
.unwrap_or(Ok(default))
};
let limit = parse_size("limit", 25)?;
let offset = parse_size("offset", 0)?;
if !(1..=100).contains(&limit) || offset > 10_000 {
return Err("Health pagination exceeds the supported range");
}
Ok(Self {
kind,
from_unix_ms: now_ms.saturating_sub(hours * 3_600_000),
to_unix_ms: now_ms,
limit,
offset,
})
}
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct HealthCoverage {
pub status: &'static str,
pub sample_status: &'static str,
pub classified_count: u64,
pub unknown_failure_count: u64,
pub excluded_count: u64,
pub exclusion_policy: &'static str,
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct PublicHealthObject {
pub id: String,
pub kind: HealthObservationObjectKind,
pub name: String,
pub status: HealthStatus,
pub request_count: u64,
pub request_success: HealthRatio,
pub service_availability: HealthRatio,
pub coverage: HealthCoverage,
pub average_latency_ms: Option<f64>,
pub latency_sample_count: u64,
pub last_request_at: Option<String>,
pub timeline: Vec<PublicHealthBucket>,
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct PublicHealthBucket {
pub from: String,
pub to: String,
pub status: HealthStatus,
pub service_availability: HealthRatio,
pub unknown_failure_count: u64,
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct AdminHealthObject {
#[serde(flatten)]
pub service: PublicHealthObject,
pub source_value: String,
pub attempts: AdminAttemptMetrics,
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct AdminAttemptMetrics {
pub succeeded_count: u64,
pub failed_count: u64,
pub in_progress_count: u64,
pub cancelled_count: u64,
pub success: HealthRatio,
}
pub(super) fn public_projection(
id: String,
kind: HealthObservationObjectKind,
name: String,
metrics: &HealthObservationMetrics,
policy: &HealthPolicy,
) -> PublicHealthObject {
let classified = metrics
.service_succeeded_count
.saturating_add(metrics.service_failed_count);
PublicHealthObject {
id,
kind,
name,
status: policy.status(
metrics.service_succeeded_count,
metrics.service_failed_count,
metrics.unknown_failure_count,
),
request_count: metrics.request_count,
request_success: HealthRatio::new(
metrics.succeeded_count,
metrics
.succeeded_count
.saturating_add(metrics.failed_count)
.saturating_add(metrics.cancelled_count),
),
service_availability: HealthRatio::new(metrics.service_succeeded_count, classified),
coverage: HealthCoverage {
status: if metrics.unknown_failure_count > 0 {
"partial"
} else {
"complete"
},
sample_status: if classified == 0 {
"empty"
} else if classified < policy.minimum_samples {
"insufficient"
} else {
"sufficient"
},
classified_count: classified,
unknown_failure_count: metrics.unknown_failure_count,
excluded_count: metrics.excluded_count,
exclusion_policy: "client_cancelled_invalid_input_identity_or_quota_policy",
},
average_latency_ms: (metrics.latency_sample_count > 0)
.then(|| metrics.latency_sum_ms as f64 / metrics.latency_sample_count as f64),
latency_sample_count: metrics.latency_sample_count,
last_request_at: metrics.last_request_at_unix_ms.and_then(unix_ms_to_rfc3339),
timeline: Vec::new(),
}
}
fn response_error(status: StatusCode, detail: &str) -> Response<Body> {
(status, Json(json!({ "detail": detail }))).into_response()
}
async fn read_publication(state: &AppState) -> Result<HealthPublication, GatewayError> {
let config = match state
.read_system_config_json_value_strong(PUBLICATION_KEY)
.await?
{
Some(value) => serde_json::from_value(value).map_err(|error| {
GatewayError::Internal(format!("Invalid health publication configuration: {error}"))
})?,
None => HealthPublication::default(),
};
config
.validate()
.map_err(|error| GatewayError::Internal(error.to_string()))?;
Ok(config)
}
pub(crate) async fn build_publication_response(
state: &AppState,
body: Option<&[u8]>,
) -> Response<Body> {
if let Some(body) = body {
let config: HealthPublication = match serde_json::from_slice(body) {
Ok(config) => config,
Err(_) => {
return response_error(
StatusCode::BAD_REQUEST,
"Invalid health publication configuration",
)
}
};
if let Err(error) = config.validate() {
return response_error(StatusCode::BAD_REQUEST, error);
}
let value = match serde_json::to_value(&config) {
Ok(value) => value,
Err(_) => {
return response_error(
StatusCode::INTERNAL_SERVER_ERROR,
"Could not encode health publication",
)
}
};
return match state
.upsert_system_config_json_value(
PUBLICATION_KEY,
&value,
Some("Public status object allowlist"),
)
.await
{
Ok(_) => Json(config).into_response(),
Err(_) => response_error(
StatusCode::SERVICE_UNAVAILABLE,
"Health publication is unavailable",
),
};
}
match read_publication(state).await {
Ok(config) => Json(config).into_response(),
Err(_) => response_error(
StatusCode::SERVICE_UNAVAILABLE,
"Health publication is unavailable",
),
}
}
pub(crate) async fn build_health_v2_response(
state: &AppState,
path: &str,
query: Option<&str>,
audience: HealthAudience,
) -> Response<Body> {
let now_ms = Utc::now().timestamp_millis().max(0) as u64;
let request = match HealthRequest::parse(query, audience != HealthAudience::Admin, now_ms) {
Ok(request) => request,
Err(detail) => return response_error(StatusCode::BAD_REQUEST, detail),
};
let prefix = match audience {
HealthAudience::Public => "/api/public/health/v2/",
HealthAudience::Authenticated => "/api/users/me/health/v2/",
HealthAudience::Admin => "/api/admin/endpoints/health/v2/",
};
let tail = path.strip_prefix(prefix).unwrap_or_default();
if tail != "summary"
&& tail != "objects"
&& !tail
.strip_prefix("objects/")
.is_some_and(|id| !id.is_empty() && !id.contains('/'))
{
return response_error(StatusCode::NOT_FOUND, "Health resource not found");
}
match build_health_payload(state, &request, tail, audience, now_ms).await {
Ok(Some(payload)) => Json(payload).into_response(),
Ok(None) => response_error(
StatusCode::NOT_FOUND,
"Health object is not published or does not exist",
),
Err(error) => {
tracing::warn!(error = %crate::error::redact_error_detail(&format!("{error:?}")), "health observation query failed");
response_error(
StatusCode::SERVICE_UNAVAILABLE,
"Health observations are temporarily unavailable",
)
}
}
}
async fn admin_objects(
state: &AppState,
kind: HealthObservationObjectKind,
) -> Result<BTreeMap<String, String>, GatewayError> {
let mut objects = BTreeMap::new();
match kind {
HealthObservationObjectKind::ApiFormat => {
let providers = state.list_provider_catalog_providers(true).await?;
let ids: Vec<_> = providers
.iter()
.map(|provider| provider.id.clone())
.collect();
for endpoint in state
.list_provider_catalog_endpoints_by_provider_ids(&ids)
.await?
{
if endpoint.is_active {
objects.insert(endpoint.api_format.clone(), endpoint.api_format);
}
}
}
HealthObservationObjectKind::Provider => {
for provider in state.list_provider_catalog_providers(false).await? {
objects.insert(provider.id, provider.name);
}
}
HealthObservationObjectKind::Model => {
let mut offset = 0;
loop {
let page = state
.list_admin_global_models(&AdminGlobalModelListQuery {
offset,
limit: 500,
..Default::default()
})
.await?;
for model in &page.items {
objects.insert(model.name.clone(), model.display_name.clone());
}
offset += page.items.len();
if offset >= page.total || page.items.is_empty() {
break;
}
if offset >= 10_000 {
return Err(GatewayError::Internal(
"Health model catalog exceeds query budget".into(),
));
}
}
}
}
Ok(objects)
}
async fn build_health_payload(
state: &AppState,
request: &HealthRequest,
tail: &str,
audience: HealthAudience,
now_ms: u64,
) -> Result<Option<serde_json::Value>, GatewayError> {
let policy = HealthPolicy::default();
let public = audience == HealthAudience::Public;
let redact_internal = audience != HealthAudience::Admin;
let publication = if public {
Some(read_publication(state).await?)
} else {
None
};
if publication.as_ref().is_some_and(|config| !config.enabled) {
return Ok(None);
}
let published: BTreeMap<_, _> = publication
.as_ref()
.map(|config| {
config
.objects
.iter()
.filter(|object| object.kind == request.kind)
.map(|object| (object.value.clone(), object))
.collect()
})
.unwrap_or_default();
let mut names = if public {
published
.iter()
.map(|(value, object)| (value.clone(), object.display_name.clone()))
.collect()
} else {
admin_objects(state, request.kind).await?
};
let observation = if public && published.is_empty() {
Default::default()
} else {
state
.data
.summarize_health_observations(&HealthObservationQuery {
from_unix_ms: request.from_unix_ms,
to_unix_ms: request.to_unix_ms,
object_kind: request.kind,
object_values: public.then(|| published.keys().cloned().collect()),
segments: 24,
})
.await
.map_err(|error| GatewayError::Internal(error.to_string()))?
};
let mut metrics_by_value = BTreeMap::new();
let mut timeline_by_value = BTreeMap::new();
for object in observation.objects {
if !public || published.contains_key(&object.object_value) {
names
.entry(object.object_value.clone())
.or_insert_with(|| object.object_value.clone());
timeline_by_value.insert(object.object_value.clone(), object.timeline);
metrics_by_value.insert(object.object_value, object.metrics);
}
}
let mut objects = Vec::new();
for (value, name) in names {
let metrics = metrics_by_value.remove(&value).unwrap_or_default();
let id = if public {
published[&value].public_id.clone()
} else {
URL_SAFE_NO_PAD.encode(value.as_bytes())
};
let mut service = public_projection(id, request.kind, name, &metrics, &policy);
service.timeline = timeline_by_value
.remove(&value)
.unwrap_or_default()
.into_iter()
.map(|bucket| PublicHealthBucket {
from: unix_ms_to_rfc3339(bucket.from_unix_ms).unwrap_or_default(),
to: unix_ms_to_rfc3339(bucket.to_unix_ms).unwrap_or_default(),
status: policy.status(
bucket.metrics.service_succeeded_count,
bucket.metrics.service_failed_count,
bucket.metrics.unknown_failure_count,
),
service_availability: HealthRatio::new(
bucket.metrics.service_succeeded_count,
bucket
.metrics
.service_succeeded_count
.saturating_add(bucket.metrics.service_failed_count),
),
unknown_failure_count: bucket.metrics.unknown_failure_count,
})
.collect();
objects.push(AdminHealthObject {
service,
source_value: value,
attempts: AdminAttemptMetrics {
succeeded_count: metrics.attempt_succeeded_count,
failed_count: metrics.attempt_failed_count,
in_progress_count: metrics.attempt_in_progress_count,
cancelled_count: metrics.attempt_cancelled_count,
success: HealthRatio::new(
metrics.attempt_succeeded_count,
metrics
.attempt_succeeded_count
.saturating_add(metrics.attempt_failed_count),
),
},
});
}
let status = overall_status(objects.iter().map(|object| object.service.status));
let meta = json!({
"schema_version": 2, "metric_version": policy.version,
"scope": { "kind": match audience { HealthAudience::Public => "published", HealthAudience::Authenticated => "authenticated", HealthAudience::Admin => "installation" }, "object_kind": request.kind },
"range": { "from": unix_ms_to_rfc3339(request.from_unix_ms), "to": unix_ms_to_rfc3339(request.to_unix_ms), "timezone": "UTC", "time_basis": "request_started_at" },
"generated_at": unix_ms_to_rfc3339(now_ms),
"data_through": observation.data_through_unix_ms.and_then(unix_ms_to_rfc3339),
"freshness": if observation.data_through_unix_ms.is_some_and(|time| now_ms.saturating_sub(time) > 120_000) { "stale" } else if observation.data_through_unix_ms.is_some() { "current" } else { "unknown" },
"policy": policy,
});
let data = if tail == "summary" {
json!({ "status": status, "object_count": objects.len(),
"healthy_count": objects.iter().filter(|object| object.service.status == HealthStatus::Healthy).count(),
"degraded_count": objects.iter().filter(|object| object.service.status == HealthStatus::Degraded).count(),
"unavailable_count": objects.iter().filter(|object| object.service.status == HealthStatus::Unavailable).count(),
"unknown_count": objects.iter().filter(|object| object.service.status == HealthStatus::Unknown).count(),
"requests": public_projection(String::new(), request.kind, String::new(), &observation.overall, &policy),
})
} else if let Some(id) = tail.strip_prefix("objects/") {
let Some(object) = objects.iter().find(|object| object.service.id == id) else {
return Ok(None);
};
if redact_internal {
serde_json::to_value(&object.service)
} else {
serde_json::to_value(object)
}
.map_err(|error| GatewayError::Internal(error.to_string()))?
} else {
let items: Vec<serde_json::Value> = objects
.iter()
.skip(request.offset)
.take(request.limit)
.map(|object| {
if redact_internal {
serde_json::to_value(&object.service)
} else {
serde_json::to_value(object)
}
})
.collect::<Result<_, _>>()
.map_err(|error| GatewayError::Internal(error.to_string()))?;
json!({ "items": items, "total": objects.len(), "limit": request.limit, "offset": request.offset })
};
Ok(Some(json!({ "meta": meta, "data": data })))
}
@@ -0,0 +1,8 @@
mod api;
mod policy;
mod publication;
pub(crate) use api::{build_health_v2_response, build_publication_response, HealthAudience};
#[cfg(test)]
mod tests;
@@ -0,0 +1,95 @@
use serde::Serialize;
pub(super) const POLICY_VERSION: &str = "service-health-v1";
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub(super) enum HealthStatus {
Healthy,
Degraded,
Unavailable,
Unknown,
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct HealthPolicy {
pub version: &'static str,
pub minimum_samples: u64,
pub healthy_threshold: f64,
pub degraded_threshold: f64,
}
impl Default for HealthPolicy {
fn default() -> Self {
Self {
version: POLICY_VERSION,
minimum_samples: 20,
healthy_threshold: 0.99,
degraded_threshold: 0.95,
}
}
}
#[derive(Clone, Debug, Serialize)]
pub(super) struct HealthRatio {
pub numerator: u64,
pub denominator: u64,
pub value: Option<f64>,
}
impl HealthRatio {
pub fn new(numerator: u64, denominator: u64) -> Self {
Self {
numerator,
denominator,
value: (denominator > 0).then(|| numerator as f64 / denominator as f64),
}
}
}
impl HealthPolicy {
pub fn status(&self, successes: u64, failures: u64, unknown: u64) -> HealthStatus {
let samples = successes.saturating_add(failures);
if samples < self.minimum_samples {
return HealthStatus::Unknown;
}
if unknown > 0 {
let best_possible_rate =
successes.saturating_add(unknown) as f64 / samples.saturating_add(unknown) as f64;
return if best_possible_rate < self.degraded_threshold {
HealthStatus::Unavailable
} else {
HealthStatus::Unknown
};
}
let ratio = successes as f64 / samples as f64;
if ratio >= self.healthy_threshold {
HealthStatus::Healthy
} else if ratio >= self.degraded_threshold {
HealthStatus::Degraded
} else {
HealthStatus::Unavailable
}
}
}
pub(super) fn overall_status(statuses: impl Iterator<Item = HealthStatus>) -> HealthStatus {
let mut result = HealthStatus::Healthy;
let mut count = 0;
for status in statuses {
count += 1;
match status {
HealthStatus::Unavailable => return HealthStatus::Unavailable,
HealthStatus::Degraded => result = HealthStatus::Degraded,
HealthStatus::Unknown if result == HealthStatus::Healthy => {
result = HealthStatus::Unknown
}
_ => {}
}
}
if count == 0 {
HealthStatus::Unknown
} else {
result
}
}
@@ -0,0 +1,60 @@
use aether_data_contracts::repository::usage::HealthObservationObjectKind;
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
pub(super) const PUBLICATION_KEY: &str = "health_publication_v1";
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(super) struct HealthPublication {
pub enabled: bool,
pub objects: Vec<PublishedHealthObject>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(super) struct PublishedHealthObject {
pub public_id: String,
pub kind: HealthObservationObjectKind,
pub value: String,
pub display_name: String,
}
impl HealthPublication {
pub fn validate(&self) -> Result<(), &'static str> {
if self.objects.len() > 200 {
return Err("At most 200 public health objects may be published");
}
let mut ids = BTreeSet::new();
let mut values = BTreeSet::new();
for object in &self.objects {
if object.public_id.is_empty()
|| object.public_id.len() > 80
|| !object
.public_id
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'-' || c == b'_')
{
return Err(
"Public IDs must contain 1 to 80 letters, digits, hyphens or underscores",
);
}
if object.kind == HealthObservationObjectKind::Provider {
return Err("Internal providers cannot be published as public health objects");
}
if object.value.trim().is_empty()
|| object.value.len() > 256
|| object.display_name.trim().is_empty()
|| object.display_name.len() > 120
{
return Err("Every public object requires a bounded source value and display name");
}
if !ids.insert(object.public_id.clone())
|| !values.insert(format!("{:?}:{}", object.kind, object.value))
{
return Err("Public IDs and health source objects must be unique");
}
}
Ok(())
}
}
@@ -0,0 +1,103 @@
use super::api::{public_projection, AdminAttemptMetrics, AdminHealthObject, HealthRequest};
use super::policy::{overall_status, HealthPolicy, HealthRatio, HealthStatus};
use super::publication::HealthPublication;
use aether_data_contracts::repository::usage::{
HealthObservationMetrics, HealthObservationObjectKind,
};
use serde_json::json;
#[test]
fn health_policy_distinguishes_unknown_and_insufficient_samples() {
let policy = HealthPolicy::default();
assert_eq!(policy.status(0, 0, 0), HealthStatus::Unknown);
assert_eq!(policy.status(19, 0, 0), HealthStatus::Unknown);
assert_eq!(policy.status(100, 0, 1), HealthStatus::Unknown);
assert_eq!(policy.status(0, 100, 1), HealthStatus::Unavailable);
assert_eq!(policy.status(99, 1, 0), HealthStatus::Healthy);
assert_eq!(policy.status(96, 4, 0), HealthStatus::Degraded);
assert_eq!(policy.status(90, 10, 0), HealthStatus::Unavailable);
assert_eq!(HealthRatio::new(0, 0).value, None);
assert_eq!(
overall_status([HealthStatus::Healthy, HealthStatus::Unknown].into_iter()),
HealthStatus::Unknown
);
}
#[test]
fn health_query_rejects_ambiguous_and_internal_public_filters() {
assert!(HealthRequest::parse(Some("provider_id=secret"), true, 30_000_000).is_err());
assert!(HealthRequest::parse(Some("kind=provider"), true, 30_000_000).is_err());
assert!(HealthRequest::parse(Some("window=6h&window=24h"), false, 30_000_000).is_err());
assert!(HealthRequest::parse(Some("limit=101"), false, 30_000_000).is_err());
let query = HealthRequest::parse(
Some("kind=model&window=1h&limit=50&offset=25"),
true,
30_000_000,
)
.unwrap();
assert_eq!(query.from_unix_ms, 26_400_000);
assert_eq!(query.to_unix_ms, 30_000_000);
assert_eq!(query.offset, 25);
}
#[test]
fn public_publication_is_explicit_and_rejects_provider_and_duplicate_sources() {
let config: HealthPublication = serde_json::from_value(json!({"enabled": true, "objects": [
{"public_id": "chat", "kind": "api_format", "value": "openai:chat", "display_name": "Chat"}
]}))
.unwrap();
assert!(config.validate().is_ok());
let mut duplicated = config.clone();
duplicated.objects.push(duplicated.objects[0].clone());
assert!(duplicated.validate().is_err());
let mut provider = config;
provider.objects[0].kind = HealthObservationObjectKind::Provider;
assert!(provider.validate().is_err());
assert!(serde_json::from_value::<HealthPublication>(
json!({"enabled": true, "objects": [], "publish_all": true})
)
.is_err());
}
#[test]
fn public_dto_cannot_serialize_internal_source_or_attempts() {
let metrics = HealthObservationMetrics {
request_count: 105,
succeeded_count: 100,
failed_count: 4,
cancelled_count: 1,
service_succeeded_count: 100,
service_failed_count: 2,
excluded_count: 2,
unknown_failure_count: 1,
..Default::default()
};
let public = public_projection(
"chat".into(),
HealthObservationObjectKind::ApiFormat,
"Chat".into(),
&metrics,
&HealthPolicy::default(),
);
assert_eq!(public.request_success.denominator, 105);
assert_eq!(public.service_availability.denominator, 102);
let admin = AdminHealthObject {
service: public.clone(),
source_value: "internal-provider-id".into(),
attempts: AdminAttemptMetrics {
succeeded_count: 100,
failed_count: 20,
in_progress_count: 1,
cancelled_count: 1,
success: HealthRatio::new(100, 120),
},
};
let public_json = serde_json::to_value(public).unwrap();
assert!(public_json.get("source_value").is_none());
assert!(public_json.get("attempts").is_none());
assert!(!public_json.to_string().contains("internal-provider-id"));
assert!(serde_json::to_value(admin)
.unwrap()
.get("attempts")
.is_some());
}
@@ -4,6 +4,7 @@ mod auth_api_key_secret;
mod catalog;
mod email_templates;
mod external_models;
pub(crate) mod health_monitor;
mod identity_oauth_provider_secret;
mod multipart;
mod normalize;
@@ -232,6 +232,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
decision.route_kind.as_deref(),
) {
(Some("endpoints_manage"), http::Method::POST, Some("create_provider_key"))
| (Some("endpoints_health"), http::Method::PUT, Some("health_v2_publication"))
| (Some("endpoints_manage"), http::Method::POST, Some("create_endpoint"))
| (Some("endpoints_manage"), http::Method::POST, Some("batch_delete_keys"))
| (Some("endpoints_manage"), http::Method::POST, Some("refresh_quota"))
@@ -341,6 +342,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
| (Some("billing_manage"), http::Method::PUT, Some("update_rule"))
| (Some("billing_manage"), http::Method::POST, Some("create_collector"))
| (Some("billing_manage"), http::Method::PUT, Some("update_collector"))
| (Some("billing_manage"), http::Method::POST, Some("create_provider_expense"))
| (Some("billing_manage"), http::Method::POST, Some("create_plan"))
| (Some("billing_manage"), http::Method::PUT, Some("update_plan"))
| (Some("billing_manage"), http::Method::PATCH, Some("set_plan_status"))
+2
View File
@@ -43,6 +43,7 @@ mod data;
mod dispatch;
mod email_delivery;
mod error;
mod execution_activity;
mod execution_runtime;
mod executor;
mod fallback_metrics;
@@ -68,6 +69,7 @@ mod provider_key_auth;
mod provider_pool_demand;
pub(crate) use aether_provider_transport as provider_transport;
mod rate_limit;
mod request_activity;
mod request_candidate_queue;
mod request_candidate_runtime;
mod request_diagnostics;
@@ -669,8 +669,24 @@ pub(super) async fn run_pending_cleanup_once(app: &AppState) -> Result<(), DataL
pub(super) async fn run_stats_hourly_aggregation_once(
data: &GatewayDataState,
) -> Result<bool, DataLayerError> {
let Some(summary) = perform_stats_hourly_aggregation_once(data).await? else {
return Ok(false);
let legacy = perform_stats_hourly_aggregation_once(data).await;
let overview_progress = match super::stats_hourly::perform_overview_rebuild_once(data).await {
Ok(progress) => progress,
Err(error) => {
warn!(event_name = "overview_rebuild_failed", error = ?error,
"overview rebuild deferred; legacy statistics remain available");
0
}
};
if overview_progress > 0 {
info!(
event_name = "overview_rebuild_progress",
progress = overview_progress,
"overview dirty projections rebuilt"
);
}
let Some(summary) = legacy? else {
return Ok(overview_progress > 0);
};
info!(
@@ -22,3 +22,19 @@ pub(super) async fn perform_stats_hourly_aggregation_once(
})
.await
}
pub(super) async fn perform_overview_rebuild_once(
data: &GatewayDataState,
) -> Result<usize, DataLayerError> {
if !data.has_stats_hourly_aggregation_backend()
|| !system_config_bool(data, "enable_stats_aggregation", true).await?
{
return Ok(0);
}
let now_utc = Utc::now();
data.rebuild_overview_buckets(&StatsHourlyAggregationInput {
target_hour_utc: stats_hourly_aggregation_target_hour(now_utc),
aggregated_at: now_utc,
})
.await
}
@@ -772,57 +772,96 @@ pub(crate) fn spawn_stats_hourly_aggregation_worker(
app,
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
|app| async move {
let data = app.data.clone();
let mut deferred_since = None;
tokio::time::sleep(STATS_AGGREGATION_STARTUP_GRACE).await;
loop {
let mut processed = 0_usize;
let mut deferred = false;
while processed < STATS_HOURLY_CATCH_UP_BURST_LIMIT {
let permit = STATS_AGGREGATION_GATE
.acquire()
.await
.expect("stats aggregation gate should remain open");
if should_defer_stats_aggregation(
&app,
// Both loops belong to the singleton lease future. Losing the
// lease/shutting down drops them together; no detached task survives.
let drain = async {
let data = app.data.clone();
let mut deferred_since = None;
let mut interval = tokio::time::interval(Duration::from_secs(10));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
interval.tick().await;
if should_defer_for_database_pressure(
&data,
"stats_hourly_aggregation",
"overview_dirty_drain",
&mut deferred_since,
) {
drop(permit);
deferred = true;
break;
continue;
}
match run_stats_hourly_aggregation_once(&data).await {
Ok(true) => {
processed += 1;
tokio::time::sleep(STATS_CATCH_UP_BUCKET_PAUSE).await;
drop(permit);
// At most ten batches per tick; normal traffic and idle
// installations stop after the first empty batch.
for _ in 0..10 {
match data.drain_overview_dirty_events(Utc::now()).await {
Ok(0) => break,
Ok(_) => tokio::task::yield_now().await,
Err(err) => {
log_maintenance_worker_failure(
"overview_dirty_drain",
"tick",
&err,
);
break;
}
}
Ok(false) => break,
Err(err) => {
log_maintenance_worker_failure(
"stats_hourly_aggregation",
"tick",
&err,
);
}
}
};
let hourly = async {
let data = app.data.clone();
let mut deferred_since = None;
tokio::time::sleep(STATS_AGGREGATION_STARTUP_GRACE).await;
loop {
let mut processed = 0_usize;
let mut deferred = false;
while processed < STATS_HOURLY_CATCH_UP_BURST_LIMIT {
let permit = STATS_AGGREGATION_GATE
.acquire()
.await
.expect("stats aggregation gate should remain open");
if should_defer_stats_aggregation(
&app,
&data,
"stats_hourly_aggregation",
&mut deferred_since,
) {
drop(permit);
deferred = true;
break;
}
match run_stats_hourly_aggregation_once(&data).await {
Ok(true) => {
processed += 1;
tokio::time::sleep(STATS_CATCH_UP_BUCKET_PAUSE).await;
drop(permit);
}
Ok(false) => break,
Err(err) => {
log_maintenance_worker_failure(
"stats_hourly_aggregation",
"tick",
&err,
);
break;
}
}
}
}
if deferred {
tokio::time::sleep(MAINTENANCE_PRESSURE_RETRY_INTERVAL).await;
continue;
}
if deferred {
tokio::time::sleep(MAINTENANCE_PRESSURE_RETRY_INTERVAL).await;
continue;
}
if processed >= STATS_HOURLY_CATCH_UP_BURST_LIMIT {
continue;
}
if processed >= STATS_HOURLY_CATCH_UP_BURST_LIMIT {
continue;
}
tokio::time::sleep(duration_until_next_stats_hourly_aggregation_run(Utc::now()))
tokio::time::sleep(
duration_until_next_stats_hourly_aggregation_run(Utc::now()),
)
.await;
}
}
};
tokio::join!(drain, hourly);
},
))
}
@@ -155,6 +155,17 @@ pub(crate) fn with_error_flow_report_context(
error_flow: Value,
) -> Option<Value> {
let mut object = report_context?.as_object()?.clone();
if !object.contains_key("analytics_failure")
&& error_flow.get("source").and_then(Value::as_str) == Some("upstream_response")
&& error_flow
.get("status_code")
.and_then(Value::as_u64)
.is_some_and(|status| status >= 400)
{
object.insert("analytics_failure".into(), json!({
"origin": "upstream", "stage": "response", "reason": "upstream_response_error", "schema_version": 1,
}));
}
object.insert("error_flow".to_string(), error_flow);
Some(Value::Object(object))
}
@@ -49,6 +49,7 @@ pub(crate) struct ProviderPoolDemandSnapshot {
pub(crate) struct ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind,
provider_key_permit: Option<RuntimeSemaphorePermit>,
observed_activity: Option<crate::execution_activity::ExecutionActivityGuard>,
released: bool,
}
@@ -73,11 +74,28 @@ enum ProviderPoolInFlightGuardKind {
}
impl ProviderPoolInFlightGuard {
fn observe_execution(
guard: Option<Self>,
activity: crate::execution_activity::ExecutionActivityGuard,
) -> Self {
// Observability remains active when provider demand tracking is disabled.
let mut guard = guard.unwrap_or(Self {
kind: ProviderPoolInFlightGuardKind::Disabled,
provider_key_permit: None,
observed_activity: None,
released: false,
});
guard.observed_activity = Some(activity);
guard
}
pub(crate) async fn release(mut self) {
self.release_inner().await;
}
async fn release_inner(&mut self) {
// Observation ends with execution, before distributed permit cleanup can wait.
self.observed_activity.take();
if self.released {
return;
}
@@ -358,6 +376,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind::Disabled,
provider_key_permit: Some(provider_key_permit),
observed_activity: None,
released: false,
}),
);
@@ -369,6 +388,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind::Disabled,
provider_key_permit: Some(provider_key_permit),
observed_activity: None,
released: false,
}),
);
@@ -381,6 +401,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
counter,
},
provider_key_permit,
observed_activity: None,
released: false,
}));
}
@@ -392,6 +413,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
counter,
},
provider_key_permit,
observed_activity: None,
released: false,
}));
}
@@ -417,6 +439,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind::Disabled,
provider_key_permit: Some(provider_key_permit),
observed_activity: None,
released: false,
}),
);
@@ -431,6 +454,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind::Disabled,
provider_key_permit: Some(provider_key_permit),
observed_activity: None,
released: false,
}),
);
@@ -454,6 +478,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
renew_handle: Some(renew_handle),
},
provider_key_permit,
observed_activity: None,
released: false,
}))
}
@@ -461,6 +486,47 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit(
pub(crate) async fn acquire_provider_pool_execution_guard(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ProviderPoolInFlightAdmission, GatewayError> {
let admission = acquire_provider_pool_execution_guard_unobserved(state, plan).await?;
let ProviderPoolInFlightAdmission::Acquired(guard) = admission else {
return Ok(admission);
};
let requested_model = report_context
.and_then(|context| context.get("model"))
.and_then(serde_json::Value::as_str);
let observation_id = execution_observation_request_id(&plan.request_id, report_context);
let activity = state.execution_activity.begin(
observation_id.as_ref(),
&plan.provider_id,
plan.provider_name.as_deref(),
requested_model,
);
Ok(ProviderPoolInFlightAdmission::Acquired(Some(
ProviderPoolInFlightGuard::observe_execution(guard, activity),
)))
}
fn execution_observation_request_id<'a>(
request_id: &'a str,
report_context: Option<&serde_json::Value>,
) -> std::borrow::Cow<'a, str> {
// Transparent Responses retries have distinct audit request IDs, but share
// one server-issued logical turn ID. Count that client request once.
report_context
.and_then(|context| context.get("websocket_logical_turn_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| std::borrow::Cow::Owned(format!("ws:{value}")))
.unwrap_or(std::borrow::Cow::Borrowed(request_id))
}
/// Long-lived audio/live sockets use the same capacity permits, but are not
/// individual requests and must not contribute to per-request RPM/concurrency.
pub(crate) async fn acquire_provider_pool_execution_guard_unobserved(
state: &AppState,
plan: &ExecutionPlan,
) -> Result<ProviderPoolInFlightAdmission, GatewayError> {
let concurrent_limit = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
@@ -751,6 +817,109 @@ mod tests {
drop(replacement);
}
#[test]
fn execution_observation_deduplicates_websocket_attempts_by_logical_turn() {
let context = serde_json::json!({"websocket_logical_turn_id": "logical-turn-1"});
assert_eq!(
execution_observation_request_id("attempt-1", Some(&context)),
execution_observation_request_id("attempt-2", Some(&context)),
);
assert_ne!(
execution_observation_request_id("logical-turn-1", None),
execution_observation_request_id("attempt-1", Some(&context)),
);
assert_eq!(
execution_observation_request_id("http-request", None),
"http-request"
);
assert_eq!(
execution_observation_request_id(
"http-request",
Some(&serde_json::json!({"websocket_logical_turn_id": " "}))
),
"http-request",
);
}
#[tokio::test]
async fn execution_activity_survives_disabled_pool_tracking_and_clears_on_release_or_drop() {
let activity = Arc::new(crate::execution_activity::ExecutionActivity::default());
// Pool mode Off (without a key limit) returns None. Execution admission
// still attaches the independent observation to a disabled wrapper.
let mut guard = ProviderPoolInFlightGuard::observe_execution(
None,
activity.begin(
"request-1",
"provider-1",
Some("Provider"),
Some("client-model"),
),
);
let snapshot = activity.snapshot();
assert_eq!(snapshot["providers"][0]["current_concurrency"], 1);
assert_eq!(snapshot["providers"][0]["requests_per_minute"], 1);
assert_eq!(snapshot["models"][0]["model"], "client-model");
guard.release_inner().await;
assert_eq!(
activity.snapshot()["providers"][0]["current_concurrency"],
0
);
assert!(guard.observed_activity.is_none());
drop(guard);
let guard = ProviderPoolInFlightGuard::observe_execution(
None,
activity.begin(
"request-2",
"provider-1",
Some("Provider"),
Some("client-model"),
),
);
assert_eq!(activity.snapshot()["models"][0]["current_concurrency"], 1);
drop(guard);
let snapshot = activity.snapshot();
assert_eq!(snapshot["models"][0]["current_concurrency"], 0);
assert_eq!(snapshot["providers"][0]["requests_per_minute"], 2);
}
#[tokio::test]
async fn execution_activity_releases_with_an_existing_pool_guard() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let activity = Arc::new(crate::execution_activity::ExecutionActivity::default());
let provider_id = "provider-observed-release";
let pool_guard = acquire_provider_pool_in_flight_guard(
runtime.clone(),
provider_id,
"request-1",
Some("candidate-1"),
"key-1",
)
.await
.expect("pool guard should be acquired");
let guard = ProviderPoolInFlightGuard::observe_execution(
Some(pool_guard),
activity.begin("request-1", provider_id, Some("Provider"), None),
);
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await,
1
);
assert_eq!(
activity.snapshot()["providers"][0]["current_concurrency"],
1
);
guard.release().await;
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await,
0
);
assert_eq!(
activity.snapshot()["providers"][0]["current_concurrency"],
0
);
}
#[tokio::test]
async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
+302
View File
@@ -0,0 +1,302 @@
//! Node-local request concurrency, integrated at lifecycle edges rather than sampled.
//!
//! Minute buckets bound memory independently of traffic volume. A request spanning
//! a bucket/day boundary contributes to both sides, including while no API polls us.
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use aether_data_contracts::repository::usage::UsageDashboardAnalyticsQuery;
use aether_runtime::{AdmissionPermit, AdmissionPermitHealth};
use chrono::{DateTime, Utc};
use serde_json::{json, Value};
const MINUTE_US: i64 = 60_000_000;
const RETAIN_MINUTES: i64 = 48 * 60;
#[derive(Debug, Default)]
struct Minute {
start_us: i64,
request_microseconds: u128,
peak: u64,
}
#[derive(Debug)]
struct History {
observed_from_us: i64,
through_us: i64,
active: u64,
minutes: VecDeque<Minute>,
}
impl History {
fn new(now_us: i64) -> Self {
Self {
observed_from_us: now_us,
through_us: now_us,
active: 0,
minutes: VecDeque::new(),
}
}
fn minute(&mut self, at_us: i64) -> &mut Minute {
let start_us = at_us.div_euclid(MINUTE_US) * MINUTE_US;
if self
.minutes
.back()
.is_none_or(|bucket| bucket.start_us != start_us)
{
self.minutes.push_back(Minute {
start_us,
..Minute::default()
});
}
self.minutes.back_mut().expect("minute was inserted")
}
fn advance(&mut self, now_us: i64) {
let now_us = now_us.max(self.through_us);
let retained_from = (now_us.div_euclid(MINUTE_US) - RETAIN_MINUTES) * MINUTE_US;
let mut cursor = self.through_us.max(retained_from);
while cursor < now_us {
let end = ((cursor.div_euclid(MINUTE_US) + 1) * MINUTE_US).min(now_us);
let active = self.active;
let bucket = self.minute(cursor);
bucket.request_microseconds += u128::from(active) * (end - cursor) as u128;
bucket.peak = bucket.peak.max(active);
cursor = end;
}
self.through_us = now_us;
while self
.minutes
.front()
.is_some_and(|bucket| bucket.start_us < retained_from)
{
self.minutes.pop_front();
}
}
fn change(&mut self, now_us: i64, entering: bool) {
self.advance(now_us);
self.active = if entering {
self.active.saturating_add(1)
} else {
self.active.saturating_sub(1)
};
let active = self.active;
let at_us = self.through_us;
let bucket = self.minute(at_us);
bucket.peak = bucket.peak.max(active);
}
fn today(&mut self, timezone: &str, now_us: i64) -> Result<Value, String> {
self.advance(now_us);
let through = DateTime::from_timestamp_micros(self.through_us)
.ok_or_else(|| "invalid concurrency observation timestamp".to_string())?;
let day_start = UsageDashboardAnalyticsQuery {
timezone: timezone.into(),
}
.today_start(through)
.map_err(|error| error.to_string())?;
let day_start_us = day_start.timestamp_micros();
// Current IANA offsets/day boundaries are minute aligned. Refuse to
// misrepresent an unsupported sub-minute historical boundary as exact.
if day_start_us.rem_euclid(MINUTE_US) != 0 {
return Err("concurrency day boundary is not minute aligned".into());
}
let observed_from_us = self.observed_from_us.max(day_start_us);
let duration_us = self.through_us.saturating_sub(observed_from_us);
let (area, peak) = self
.minutes
.iter()
.filter(|minute| minute.start_us >= day_start_us && minute.start_us <= self.through_us)
.fold((0u128, self.active), |(area, peak), minute| {
(area + minute.request_microseconds, peak.max(minute.peak))
});
Ok(json!({
"avg": (duration_us > 0).then(|| area as f64 / duration_us as f64),
"peak": peak,
"observed_from": DateTime::from_timestamp_micros(observed_from_us),
"observed_through": through,
"scope": "node",
"measurement": "http_and_responses_websocket_requests",
"coverage": if self.observed_from_us <= day_start_us { "complete" } else { "partial" },
}))
}
}
#[derive(Debug)]
pub(crate) struct RequestActivity {
started_at: Instant,
started_at_us: i64,
history: Mutex<History>,
}
impl Default for RequestActivity {
fn default() -> Self {
let started_at = Instant::now();
let started_at_us = Utc::now().timestamp_micros();
Self {
started_at,
started_at_us,
history: Mutex::new(History::new(started_at_us)),
}
}
}
impl RequestActivity {
fn now_us(&self) -> i64 {
self.started_at_us
.saturating_add(self.started_at.elapsed().as_micros().min(i64::MAX as u128) as i64)
}
pub(crate) fn begin(self: &Arc<Self>) -> RequestActivityGuard {
self.history
.lock()
.unwrap_or_else(|error| error.into_inner())
.change(self.now_us(), true);
RequestActivityGuard {
activity: Arc::clone(self),
}
}
pub(crate) fn today(&self, timezone: &str) -> Result<Value, String> {
self.history
.lock()
.unwrap_or_else(|error| error.into_inner())
.today(timezone, self.now_us())
}
#[cfg(test)]
pub(crate) fn active(&self) -> u64 {
self.history.lock().unwrap().active
}
}
#[derive(Debug)]
pub(crate) struct RequestActivityGuard {
activity: Arc<RequestActivity>,
}
impl RequestActivityGuard {
pub(crate) fn into_admission_permit(self) -> AdmissionPermit {
// This guard observes lifecycle only; it neither limits nor cancels work.
AdmissionPermit::from_parts(None, Some(self)).expect("activity guard is present")
}
}
impl AdmissionPermitHealth for RequestActivityGuard {
fn is_healthy(&self) -> bool {
true
}
fn requires_health_poll(&self) -> bool {
false
}
}
impl Drop for RequestActivityGuard {
fn drop(&mut self) {
self.activity
.history
.lock()
.unwrap_or_else(|error| error.into_inner())
.change(self.activity.now_us(), false);
}
}
impl crate::AppState {
pub(crate) fn today_concurrency(&self, timezone: &str) -> Result<Value, String> {
self.request_activity.today(timezone)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn at(value: &str) -> i64 {
DateTime::parse_from_rfc3339(value)
.unwrap()
.timestamp_micros()
}
#[test]
fn concurrency_integrates_time_instead_of_averaging_event_samples() {
let start = at("2026-09-19T00:00:00Z");
let mut history = History::new(start);
history.change(start, true);
history.change(start + 10_000_000, true);
history.change(start + 20_000_000, false);
history.change(start + 30_000_000, false);
let value = history.today("UTC", start + 100_000_000).unwrap();
assert_eq!(value["avg"], 0.4);
assert_eq!(value["peak"], 2);
assert_eq!(value["coverage"], "complete");
}
#[test]
fn concurrency_long_request_crosses_minutes_and_local_midnight() {
let start = at("2026-09-18T15:59:30Z");
let mut history = History::new(start);
history.change(start, true);
let value = history.today("Asia/Shanghai", start + 150_000_000).unwrap();
assert_eq!(value["avg"], 1.0);
assert_eq!(value["peak"], 1);
assert_eq!(value["observed_from"], "2026-09-18T16:00:00Z");
assert_eq!(value["coverage"], "complete");
history.change(start + 150_000_000, false);
let value = history.today("Asia/Shanghai", start + 270_000_000).unwrap();
assert_eq!(value["avg"], 0.5);
}
#[test]
fn concurrency_restart_only_claims_the_observed_part_of_the_day() {
let start = at("2026-09-19T12:00:00Z");
let mut history = History::new(start);
let value = history.today("UTC", start + 60_000_000).unwrap();
assert_eq!(value["avg"], 0.0);
assert_eq!(value["peak"], 0);
assert_eq!(value["coverage"], "partial");
assert_eq!(value["observed_from"], "2026-09-19T12:00:00Z");
assert!(History::new(start).today("UTC", start).unwrap()["avg"].is_null());
assert!(history.today("not/a/timezone", start).is_err());
}
#[test]
fn concurrency_does_not_carry_yesterdays_peak_into_today() {
let start = at("2026-09-18T23:59:30Z");
let mut history = History::new(start);
history.change(start, true);
history.change(start, true);
history.change(start + 20_000_000, false);
history.change(start + 30_000_000, false);
let value = history.today("UTC", start + 90_000_000).unwrap();
assert_eq!(value["avg"], 0.0);
assert_eq!(value["peak"], 0);
}
#[test]
fn concurrency_handles_dst_and_bounds_memory_after_a_long_idle_gap() {
let start = at("2026-10-30T00:00:00Z");
let mut history = History::new(start);
history.change(start, true);
let end = at("2026-11-02T04:30:00Z");
let value = history.today("America/New_York", end).unwrap();
assert_eq!(value["avg"], 1.0);
assert_eq!(value["observed_from"], "2026-11-01T04:00:00Z");
assert_eq!(value["coverage"], "complete");
assert!(history.minutes.len() <= RETAIN_MINUTES as usize + 1);
}
#[test]
fn concurrency_permit_clones_share_one_lifecycle() {
let activity = Arc::new(RequestActivity::default());
let permit = activity.begin().into_admission_permit();
let background = permit.clone();
assert_eq!(activity.active(), 1);
drop(permit);
assert_eq!(activity.active(), 1);
drop(background);
assert_eq!(activity.active(), 0);
}
}
+88 -23
View File
@@ -1,7 +1,7 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use aether_routing_core::RoutingExecutionPolicy;
@@ -16,6 +16,15 @@ use crate::GatewayError;
tokio::task_local! {
static CANCEL_ON_CLIENT_DISCONNECT: Arc<AtomicBool>;
static REQUEST_ACTIVITY: Arc<Mutex<Option<aether_runtime::AdmissionPermit>>>;
}
/// Begin observing only after routing chose local AI execution. The surrounding
/// request/body lifecycle keeps this holder alive through disconnect draining.
pub(crate) fn track_request_activity(permit: aether_runtime::AdmissionPermit) {
let _ = REQUEST_ACTIVITY.try_with(|activity| {
*activity.lock().unwrap_or_else(|error| error.into_inner()) = Some(permit);
});
}
pub(crate) fn configure_client_disconnect(policy: RoutingExecutionPolicy) {
@@ -59,23 +68,35 @@ where
let diagnostics = Arc::new(RequestDiagnostics::default());
let cancel_for_response = Arc::clone(&cancel);
let producer_for_request = producer.clone();
let future = CANCEL_ON_CLIENT_DISCONNECT.scope(
Arc::clone(&cancel),
scope_request_diagnostics_with(Some(Arc::clone(&diagnostics)), async move {
let response = future.await?;
let complete_on_disconnect = !cancel_for_response.load(Ordering::Acquire);
if !complete_on_disconnect && producer.is_none() {
return Ok(response);
}
Ok(response.map(|body| {
Body::new(CompleteOnDisconnectBody {
body: Some(body),
diagnostics,
complete_on_disconnect,
producer,
})
}))
}),
let activity = Arc::new(Mutex::new(None));
let activity_for_response = Arc::clone(&activity);
let future = REQUEST_ACTIVITY.scope(
activity,
CANCEL_ON_CLIENT_DISCONNECT.scope(
Arc::clone(&cancel),
scope_request_diagnostics_with(Some(Arc::clone(&diagnostics)), async move {
let response = future.await?;
let complete_on_disconnect = !cancel_for_response.load(Ordering::Acquire);
if !complete_on_disconnect
&& producer.is_none()
&& activity_for_response
.lock()
.unwrap_or_else(|error| error.into_inner())
.is_none()
{
return Ok(response);
}
Ok(response.map(|body| {
Body::new(CompleteOnDisconnectBody {
body: Some(body),
diagnostics,
complete_on_disconnect,
producer,
activity: Some(activity_for_response),
})
}))
}),
),
);
CompleteOnDisconnectRequest {
future: Some(Box::pin(future)),
@@ -142,6 +163,7 @@ struct CompleteOnDisconnectBody {
complete_on_disconnect: bool,
// Drop the body first so its terminal handoff registers before this guard ends.
producer: Option<Arc<UsageProducerGuard>>,
activity: Option<Arc<Mutex<Option<aether_runtime::AdmissionPermit>>>>,
}
impl HttpBody for CompleteOnDisconnectBody {
@@ -159,6 +181,7 @@ impl HttpBody for CompleteOnDisconnectBody {
if matches!(result, Poll::Ready(None | Some(Err(_)))) {
self.body.take();
self.producer.take();
self.activity.take();
}
result
}
@@ -185,10 +208,12 @@ impl Drop for CompleteOnDisconnectBody {
};
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
let producer = self.producer.take();
let activity = self.activity.take();
runtime.spawn(scope_request_diagnostics_with(
Some(Arc::clone(&self.diagnostics)),
async move {
let _producer = producer;
let _activity = activity;
drain_body(body).await;
},
));
@@ -341,10 +366,13 @@ mod tests {
#[tokio::test]
async fn usage_shutdown_waits_for_a_disconnected_request_before_headers() {
let usage = Arc::new(UsageRuntime::disabled());
let activity = Arc::new(crate::request_activity::RequestActivity::default());
let activity_permit = activity.begin().into_admission_permit();
let (started_tx, started_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel::<()>();
let request = tokio::spawn(run_request_with_usage(usage.clone(), async move {
configure_client_disconnect(RoutingExecutionPolicy::default());
track_request_activity(activity_permit);
started_tx.send(()).unwrap();
release_rx.await.unwrap();
Ok(Response::new(Body::empty()))
@@ -354,49 +382,79 @@ mod tests {
assert!(request.await.unwrap_err().is_cancelled());
assert!(usage.shutdown(Duration::from_millis(30)).await.is_err());
assert_eq!(usage.metrics_snapshot().producers_in_flight, 1);
assert_eq!(activity.active(), 1);
release_tx.send(()).unwrap();
usage.shutdown(Duration::from_secs(1)).await.unwrap();
assert_eq!(usage.metrics_snapshot().producers_in_flight, 0);
assert_eq!(activity.active(), 0);
}
#[tokio::test]
async fn request_activity_releases_when_the_handler_fails_before_headers() {
let activity = Arc::new(crate::request_activity::RequestActivity::default());
let permit = activity.begin().into_admission_permit();
let result = run_request(async move {
track_request_activity(permit);
Err(GatewayError::Internal("test failure".into()))
})
.await;
assert!(result.is_err());
assert_eq!(activity.active(), 0);
}
#[tokio::test]
async fn usage_shutdown_waits_for_disconnected_body_drain() {
let usage = Arc::new(UsageRuntime::disabled());
let activity = Arc::new(crate::request_activity::RequestActivity::default());
let activity_permit = activity.begin().into_admission_permit();
let (sender, receiver) = mpsc::channel::<Result<Bytes, io::Error>>(1);
let response = run_request_with_usage(usage.clone(), async move {
configure_client_disconnect(RoutingExecutionPolicy::default());
Ok(Response::new(Body::from_stream(stream::unfold(
track_request_activity(activity_permit);
let response = Response::new(Body::from_stream(stream::unfold(
receiver,
|mut receiver| async { receiver.recv().await.map(|item| (item, receiver)) },
))))
)));
Ok(response)
})
.await
.unwrap();
drop(response);
assert_eq!(
activity.active(),
1,
"background drain still owns the request"
);
assert!(usage.shutdown(Duration::from_millis(30)).await.is_err());
sender.send(Ok(Bytes::from_static(b"last"))).await.unwrap();
drop(sender);
usage.shutdown(Duration::from_secs(1)).await.unwrap();
assert_eq!(usage.metrics_snapshot().producers_in_flight, 0);
assert_eq!(activity.active(), 0);
}
#[tokio::test]
async fn tracked_bodies_release_shutdown_on_cancellation_or_eof() {
for cancel_on_client_disconnect in [false, true] {
let usage = Arc::new(UsageRuntime::disabled());
let activity = Arc::new(crate::request_activity::RequestActivity::default());
let activity_permit = activity.begin().into_admission_permit();
let (sender, receiver) = mpsc::channel::<Result<Bytes, io::Error>>(1);
let response = run_request_with_usage(usage.clone(), async move {
configure_client_disconnect(RoutingExecutionPolicy {
cancel_on_client_disconnect,
..Default::default()
});
Ok(Response::new(Body::from_stream(stream::unfold(
track_request_activity(activity_permit);
let response = Response::new(Body::from_stream(stream::unfold(
receiver,
|mut receiver| async { receiver.recv().await.map(|item| (item, receiver)) },
))))
)));
Ok(response)
})
.await
.unwrap();
assert_eq!(activity.active(), 1);
let mut body = response.into_body();
if cancel_on_client_disconnect {
drop(body);
@@ -407,13 +465,17 @@ mod tests {
assert_eq!(usage.metrics_snapshot().producers_in_flight, 0);
}
usage.shutdown(Duration::from_secs(1)).await.unwrap();
assert_eq!(activity.active(), 0);
}
}
#[tokio::test]
async fn connected_response_preserves_headers_size_hint_and_trailers() {
let response = run_request(async {
let activity = Arc::new(crate::request_activity::RequestActivity::default());
let activity_permit = activity.begin().into_admission_permit();
let response = run_request(async move {
configure_client_disconnect(RoutingExecutionPolicy::default());
track_request_activity(activity_permit);
Ok(Response::builder()
.status(201)
.header("x-test", "unchanged")
@@ -429,11 +491,14 @@ mod tests {
response.into_body().collect().await.unwrap().to_bytes(),
"hello"
);
assert_eq!(activity.active(), 0);
let mut trailers = HeaderMap::new();
trailers.insert("x-finished", "yes".parse().unwrap());
let activity_permit = activity.begin().into_admission_permit();
let response = run_request(async move {
configure_client_disconnect(RoutingExecutionPolicy::default());
track_request_activity(activity_permit);
let frames = stream::iter([
Ok::<_, io::Error>(Frame::data(Bytes::from_static(b"hello"))),
Ok(Frame::trailers(trailers)),
+4 -1
View File
@@ -18,7 +18,7 @@ use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
AuthApiKeyFeatureCacheKey, AuthApiKeyIdentityCacheKey, AuthApiKeyLastUsedCache,
AuthContextCache, AuthSnapshotCache, DashboardResponseCache, DirectPlanBypassCache,
JsonValueCache, SchedulerAffinityCache, SystemConfigCache, ValueCache,
JsonValueCache, OverviewTotalCache, SchedulerAffinityCache, SystemConfigCache, ValueCache,
};
use super::super::data::GatewayDataState;
use super::super::fallback_metrics;
@@ -388,6 +388,8 @@ pub struct AppState {
pub(crate) runtime_state: Arc<RuntimeState>,
pub(crate) internal_gateway_auth: Arc<crate::internal_gateway_auth::InternalGatewayAuthConfig>,
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
pub(crate) request_activity: Arc<crate::request_activity::RequestActivity>,
pub(crate) execution_activity: Arc<crate::execution_activity::ExecutionActivity>,
pub(crate) video_tasks: Arc<VideoTaskService>,
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
@@ -429,6 +431,7 @@ pub struct AppState {
pub(crate) scheduler_affinity_cache: Arc<SchedulerAffinityCache>,
pub(crate) scheduler_affinity_epoch: Arc<AtomicU64>,
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
pub(crate) overview_total_cache: Arc<OverviewTotalCache>,
pub(crate) system_config_cache: Arc<SystemConfigCache>,
pub(crate) endpoint_response_header_rules_cache: Arc<JsonValueCache<String>>,
pub(crate) candidate_row_page_cache: Arc<super::super::cache::CandidateRowPageCache>,
+7 -2
View File
@@ -38,8 +38,9 @@ use super::super::async_task::{
};
use super::super::cache::{
AuthApiKeyLastUsedCache, AuthContextCache, AuthSnapshotCache, DashboardResponseCache,
DirectPlanBypassCache, JsonValueCache, SchedulerAffinityCache, SchedulerAffinitySnapshotEntry,
SchedulerAffinityTarget, SystemConfigCache, SystemConfigInflightRegistration, ValueCache,
DirectPlanBypassCache, JsonValueCache, OverviewTotalCache, SchedulerAffinityCache,
SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget, SystemConfigCache,
SystemConfigInflightRegistration, ValueCache,
};
use super::super::data::{GatewayDataConfig, GatewayDataState};
use super::super::fallback_metrics;
@@ -254,6 +255,7 @@ impl AppState {
}
fn replace_foreground_data_state(&mut self, data: Arc<GatewayDataState>) {
self.overview_total_cache = Arc::new(OverviewTotalCache::default());
self.clear_provider_transport_snapshot_cache();
self.invalidate_scheduler_affinity_cache();
self.invalidate_auth_context_cache();
@@ -351,6 +353,8 @@ impl AppState {
runtime_state: runtime_state.clone(),
internal_gateway_auth,
usage_runtime: Arc::new(usage::UsageRuntime::disabled()),
request_activity: Arc::new(crate::request_activity::RequestActivity::default()),
execution_activity: Arc::new(crate::execution_activity::ExecutionActivity::default()),
video_tasks: Arc::new(VideoTaskService::new(
VideoTaskTruthSourceMode::PythonSyncReport,
)),
@@ -402,6 +406,7 @@ impl AppState {
scheduler_affinity_cache: Arc::new(SchedulerAffinityCache::default()),
scheduler_affinity_epoch: Arc::new(AtomicU64::new(0)),
dashboard_response_cache: Arc::new(DashboardResponseCache::default()),
overview_total_cache: Arc::new(OverviewTotalCache::default()),
system_config_cache: Arc::new(SystemConfigCache::default()),
endpoint_response_header_rules_cache: Arc::new(JsonValueCache::default()),
candidate_row_page_cache: Arc::new(crate::cache::CandidateRowPageCache::default()),
@@ -1,6 +1,18 @@
use crate::{AppState, GatewayError};
impl AppState {
pub(crate) async fn list_user_announcements(
&self,
user_id: &str,
query: &aether_data::repository::announcements::UserAnnouncementListQuery,
) -> Result<aether_data::repository::announcements::StoredUserAnnouncementPage, GatewayError>
{
self.data
.list_user_announcements(user_id, query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_announcements(
&self,
query: &aether_data::repository::announcements::AnnouncementListQuery,
@@ -610,6 +610,17 @@ impl AppState {
.map_err(data_error)
}
pub(crate) async fn list_user_plan_entitlements_with_history(
&self,
user_id: &str,
include_inactive: bool,
) -> Result<Option<Vec<UserPlanEntitlementRecord>>, GatewayError> {
self.data
.list_user_plan_entitlements_with_history(user_id, include_inactive)
.await
.map_err(data_error)
}
pub(crate) async fn revoke_user_plan_entitlement(
&self,
user_id: &str,
@@ -14,6 +14,7 @@ use crate::{
impl AppState {
pub(crate) async fn list_admin_wallets(
&self,
user_id: Option<&str>,
status: Option<&str>,
owner_type: Option<&str>,
limit: usize,
@@ -22,6 +23,7 @@ impl AppState {
let page = self
.data
.list_admin_wallets(&AdminWalletListQuery {
user_id: user_id.map(ToOwned::to_owned),
status: status.map(ToOwned::to_owned),
owner_type: owner_type.map(ToOwned::to_owned),
limit,
@@ -3,6 +3,35 @@ use aether_data_contracts::repository::{candidates, usage};
use usage::{StoredUsageDailySummary, UsageDailyHeatmapQuery};
impl AppState {
pub(crate) async fn query_dashboard_analytics(
&self,
query: &usage::UsageDashboardAnalyticsQuery,
) -> Result<usage::StoredUsageDashboardAnalytics, GatewayError> {
self.data
.query_dashboard_analytics(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn query_usage_analytics(
&self,
query: &usage::UsageAnalyticsQuery,
) -> Result<usage::StoredUsageAnalytics, GatewayError> {
self.data
.query_usage_analytics(query)
.await
.map_err(|err| match err {
aether_data_contracts::DataLayerError::InvalidInput(message)
if query.view == usage::UsageAnalyticsView::DashboardCharts =>
{
GatewayError::Client {
status: http::StatusCode::UNPROCESSABLE_ENTITY,
message,
}
}
err => GatewayError::Internal(err.to_string()),
})
}
#[allow(dead_code)]
pub(crate) async fn rebuild_api_key_usage_stats(&self) -> Result<u64, GatewayError> {
self.data
@@ -616,6 +616,7 @@ async fn gateway_handles_admin_api_keys_create_locally_with_trusted_admin_princi
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["name"], json!("standalone-key"));
assert_eq!(payload["is_standalone"], json!(true));
assert!(payload.get("credential_kind").is_none());
assert_eq!(payload["rate_limit"], serde_json::Value::Null);
assert_eq!(payload["concurrent_limit"], serde_json::Value::Null);
assert_eq!(payload["allowed_providers"], json!(["openai"]));
@@ -646,6 +647,7 @@ async fn gateway_handles_admin_api_keys_create_locally_with_trusted_admin_princi
list_response.json().await.expect("list json should parse");
assert_eq!(list_payload["total"], json!(1));
assert_eq!(list_payload["api_keys"][0]["name"], json!("standalone-key"));
assert!(list_payload["api_keys"][0].get("credential_kind").is_none());
gateway_handle.abort();
upstream_handle.abort();
@@ -702,6 +704,7 @@ async fn gateway_handles_admin_api_keys_update_locally_with_trusted_admin_princi
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["id"], json!("key-123"));
assert_eq!(payload["name"], json!("renamed-key"));
assert!(payload.get("credential_kind").is_none());
assert_eq!(payload["rate_limit"], serde_json::Value::Null);
assert_eq!(payload["concurrent_limit"], json!(12));
assert_eq!(payload["allowed_providers"], json!(["gemini"]));
@@ -72,6 +72,138 @@ async fn send_admin_billing_request(
request.send().await.expect("request should succeed")
}
#[tokio::test]
async fn user_account_history_http_filters_wallet_and_plan_history() {
use aether_data::repository::{
billing::{InMemoryBillingReadRepository, UserPlanEntitlementRecord},
users::StoredUserAuthRecord,
wallet::{InMemoryWalletRepository, StoredWalletSnapshot},
};
let users = ["user-1", "user-2"].map(|id| {
StoredUserAuthRecord::new(
id.to_string(),
Some(format!("{id}@example.com")),
true,
id.to_string(),
Some("hash".to_string()),
"user".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
None,
None,
)
.expect("user should build")
});
let wallets = ["user-1", "user-2"].map(|id| {
StoredWalletSnapshot::new(
format!("wallet-{id}"),
Some(id.to_string()),
None,
12.5,
2.5,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
30.0,
10.0,
3.0,
1.5,
1_710_000_000,
)
.expect("wallet should build")
});
let now = chrono::Utc::now().timestamp().max(0) as u64;
let entitlements = [
("current", "user-1", "active"),
("revoked", "user-1", "revoked"),
("another-user", "user-2", "revoked"),
]
.map(|(id, user_id, status)| UserPlanEntitlementRecord {
id: id.to_string(),
user_id: user_id.to_string(),
plan_id: "plan-1".to_string(),
payment_order_id: format!("order-{id}"),
status: status.to_string(),
starts_at_unix_secs: now - 60,
expires_at_unix_secs: now + 3600,
entitlements_snapshot: json!([]),
created_at_unix_secs: now - 60,
updated_at_unix_secs: now,
});
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::with_user_billing_and_wallet_for_tests(
Arc::new(InMemoryUserReadRepository::seed_auth_users(users)),
Arc::new(InMemoryBillingReadRepository::seed_user_plan_entitlements(
entitlements,
)),
Arc::new(InMemoryWalletRepository::seed(wallets)),
),
);
let (url, handle) = start_server(build_router_with_state(state)).await;
let wallet_path = "/api/admin/wallets?user_id=user-1&owner_type=user&limit=1&offset=0";
let response = send_admin_billing_request(&url, http::Method::GET, wallet_path, None).await;
assert_eq!(response.status(), StatusCode::OK);
let wallet: serde_json::Value = response.json().await.unwrap();
assert_eq!(wallet["total"], 1);
assert_eq!(wallet["items"].as_array().unwrap().len(), 1);
assert_eq!(wallet["items"][0]["id"], "wallet-user-1");
assert_eq!(wallet["items"][0]["user_id"], "user-1");
let response = send_admin_billing_request(
&url,
http::Method::GET,
"/api/admin/wallets?user_id=missing-user&limit=1",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let missing: serde_json::Value = response.json().await.unwrap();
assert_eq!(missing["total"], 0);
assert_eq!(missing["items"], json!([]));
let path = "/api/admin/users/user-1/billing/entitlements";
for query in ["", "?include_inactive=false"] {
let response =
send_admin_billing_request(&url, http::Method::GET, &format!("{path}{query}"), None)
.await;
assert_eq!(response.status(), StatusCode::OK);
let current: serde_json::Value = response.json().await.unwrap();
assert_eq!(current["total"], 1);
assert_eq!(current["items"][0]["id"], "current");
assert_eq!(current["items"][0]["active"], true);
}
let response = send_admin_billing_request(
&url,
http::Method::GET,
&format!("{path}?include_inactive=true"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let history: serde_json::Value = response.json().await.unwrap();
assert_eq!(history["total"], 2);
let items = history["items"].as_array().unwrap();
assert!(items.iter().all(|item| item["user_id"] == "user-1"));
assert!(items
.iter()
.any(|item| item["id"] == "revoked" && item["active"] == false));
let response = send_admin_billing_request(
&url,
http::Method::GET,
&format!("{path}?include_inactive=invalid"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_billing_presets_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -657,3 +789,110 @@ async fn gateway_handles_admin_billing_collector_routes_locally_with_trusted_adm
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn provider_expense_http_contract_requires_admin_and_records_retries_only_once() {
use aether_data::repository::{
billing::InMemoryBillingReadRepository,
provider_catalog::InMemoryProviderCatalogReadRepository,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
let provider = StoredProviderCatalogProvider::new(
"provider-1".into(),
"=Supplier".into(),
None,
"custom".into(),
)
.unwrap();
let catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::with_billing_reader_for_tests(Arc::new(
InMemoryBillingReadRepository::default(),
))
.with_provider_catalog_reader(catalog),
);
let (url, handle) = start_server(build_router_with_state(state)).await;
let path = "/api/admin/billing/provider-expenses";
let client = reqwest::Client::new();
assert!(matches!(
client
.get(format!("{url}{path}"))
.send()
.await
.unwrap()
.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let payload = json!({"client_request_id":uuid::Uuid::new_v4().to_string(),"provider_id":"provider-1","kind":"subscription","amount":"12.30","currency":"USD","paid_at":"2026-09-20T00:00:00Z","period_start":"2026-09-20T00:00:00Z","period_end":"2026-10-20T00:00:00Z","note":"=SUM(1,2)"});
let first =
send_admin_billing_request(&url, http::Method::POST, path, Some(payload.clone())).await;
assert_eq!(first.status(), StatusCode::OK);
let first: serde_json::Value = first.json().await.unwrap();
assert_eq!(first["item"]["amount"], "12.30000000");
let again =
send_admin_billing_request(&url, http::Method::POST, path, Some(payload.clone())).await;
assert_eq!(again.status(), StatusCode::OK);
let again: serde_json::Value = again.json().await.unwrap();
assert_eq!(first["item"]["id"], again["item"]["id"]);
let mut conflict = payload.clone();
conflict["amount"] = json!("13");
assert_eq!(
send_admin_billing_request(&url, http::Method::POST, path, Some(conflict))
.await
.status(),
StatusCode::CONFLICT
);
let range = "?from=2026-09-01T00:00:00Z&to=2026-10-01T00:00:00Z&limit=1&offset=5";
let page =
send_admin_billing_request(&url, http::Method::GET, &format!("{path}{range}"), None).await;
assert_eq!(page.status(), StatusCode::OK);
let page: serde_json::Value = page.json().await.unwrap();
assert_eq!(page["total"], 1);
assert_eq!(page["items"], json!([]));
assert_eq!(page["totals"][0]["subscription_amount"], "12.30000000");
let csv = send_admin_billing_request(
&url,
http::Method::GET,
&format!("{path}{range}&format=csv"),
None,
)
.await;
assert_eq!(csv.status(), StatusCode::OK);
let csv = csv.text().await.unwrap();
assert!(csv.contains("'=Supplier"));
assert!(csv.contains("'=SUM(1,2)"));
assert!(csv.contains("12.30000000"));
let accounts = send_admin_billing_request(
&url,
http::Method::GET,
"/api/admin/billing/provider-accounts",
None,
)
.await;
assert_eq!(accounts.status(), StatusCode::OK);
let accounts: serde_json::Value = accounts.json().await.unwrap();
assert_eq!(accounts["items"][0]["provider_id"], "provider-1");
assert!(accounts["items"][0]["balance"].is_null());
let void_path = format!("{path}/{}/void", first["item"]["id"].as_str().unwrap());
let voided = send_admin_billing_request(&url, http::Method::POST, &void_path, None).await;
assert_eq!(voided.status(), StatusCode::OK);
let voided: serde_json::Value = voided.json().await.unwrap();
assert_eq!(voided["item"]["status"], "void");
let again = send_admin_billing_request(&url, http::Method::POST, &void_path, None).await;
assert_eq!(again.status(), StatusCode::OK);
let again: serde_json::Value = again.json().await.unwrap();
assert_eq!(voided, again);
let page: serde_json::Value =
send_admin_billing_request(&url, http::Method::GET, &format!("{path}{range}"), None)
.await
.json()
.await
.unwrap();
assert_eq!(page["total"], 0);
assert_eq!(page["totals"], json!([]));
handle.abort();
}
@@ -32,6 +32,107 @@ use crate::data::GatewayDataState;
const ADMIN_ENDPOINT_HEALTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin endpoint health data unavailable";
#[tokio::test]
async fn health_v2_publication_requires_admin_and_public_projection_keeps_empty_objects() {
use aether_data::repository::usage::InMemoryUsageReadRepository;
let data = GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed(Vec::new()),
))
.with_system_config_values_for_tests(vec![(
"health_publication_v1".to_string(),
json!({ "enabled": false, "objects": [] }),
)]);
let gateway = build_router_with_state(AppState::new().unwrap().with_data_state_for_tests(data));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let publication = json!({ "enabled": true, "objects": [
{"public_id": "chat", "kind": "api_format", "value": "internal-format", "display_name": "Chat API"},
{"public_id": "model", "kind": "model", "value": "internal-model", "display_name": "Model API"}
]});
let denied = client
.put(format!(
"{gateway_url}/api/admin/endpoints/health/v2/publication"
))
.json(&publication)
.send()
.await
.unwrap();
assert!(matches!(
denied.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let disabled = client
.get(format!("{gateway_url}/api/public/health/v2/objects"))
.send()
.await
.unwrap();
assert_eq!(disabled.status(), StatusCode::NOT_FOUND);
let saved = client
.put(format!(
"{gateway_url}/api/admin/endpoints/health/v2/publication"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&publication)
.send()
.await
.unwrap();
assert_eq!(saved.status(), StatusCode::OK);
assert_eq!(
saved.json::<serde_json::Value>().await.unwrap(),
publication
);
let public = client
.get(format!(
"{gateway_url}/api/public/health/v2/objects?kind=api_format&window=1h"
))
.send()
.await
.unwrap();
assert_eq!(public.status(), StatusCode::OK);
let body: serde_json::Value = public.json().await.unwrap();
assert_eq!(body["data"]["total"], 1);
assert_eq!(body["data"]["items"][0]["id"], "chat");
assert_eq!(body["data"]["items"][0]["status"], "unknown");
assert_eq!(body["data"]["items"][0]["request_count"], 0);
assert!(body["data"]["items"][0]["service_availability"]["value"].is_null());
let text = body.to_string();
for forbidden in [
"internal-format",
"internal-model",
"provider_id",
"source_value",
"attempts",
] {
assert!(
!text.contains(forbidden),
"public projection leaked {forbidden}"
);
}
let hidden = client
.get(format!(
"{gateway_url}/api/public/health/v2/objects/internal-format"
))
.send()
.await
.unwrap();
assert_eq!(hidden.status(), StatusCode::NOT_FOUND);
let internal_kind = client
.get(format!(
"{gateway_url}/api/public/health/v2/objects?kind=provider"
))
.send()
.await
.unwrap();
assert_eq!(internal_kind.status(), StatusCode::BAD_REQUEST);
gateway_handle.abort();
}
async fn assert_admin_modules_status_with_smtp_password(
stored_password: &str,
notification_ready: bool,
@@ -30,6 +30,8 @@ use crate::constants::{
};
use crate::data::GatewayDataState;
mod overview;
const DAY_0_UNIX_SECS: i64 = 1_710_913_600;
const DAY_1_UNIX_SECS: i64 = 1_711_000_000;
const DAY_2_UNIX_SECS: i64 = 1_711_086_400;
@@ -0,0 +1,461 @@
use super::*;
const RANGE: &str = "from=2026-09-01T23:45:00Z&to=2026-09-02T00:15:00Z&timezone=Asia%2FShanghai";
#[tokio::test]
async fn overview_live_rates_use_only_the_last_sixty_seconds() {
let now = chrono::Utc::now().timestamp();
let recent = sample_usage_row(
"live-recent",
"live-recent-request",
None,
None,
None,
"Provider",
"model-1",
100,
20,
1.0,
0.8,
now - 10,
);
let historical = sample_usage_row(
"live-earlier",
"live-earlier-request",
None,
None,
None,
"Provider",
"model-1",
1000,
200,
10.0,
8.0,
now - 120,
);
let repository = Arc::new(InMemoryUsageReadRepository::seed([recent, historical]));
let gateway = build_router_with_state(
AppState::new()
.unwrap()
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(repository)),
);
let (url, handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().get(format!("{url}/api/admin/overview/operations/live")),
)
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let value: serde_json::Value = response.json().await.unwrap();
let activity = &value["data"]["recent_activity"]["data"];
assert_eq!(activity["window_seconds"], 60);
assert_eq!(activity["requests_per_minute"], 1);
assert_eq!(activity["tokens_per_minute"], 120);
assert_eq!(activity["requests_per_second"], 1.0 / 60.0);
handle.abort();
}
#[tokio::test]
async fn overview_dashboard_summary_reads_only_the_new_collection_and_all_current_users() {
let now = chrono::Utc::now();
let since = now - chrono::Duration::minutes(1);
let mut recent = sample_usage_row(
"dashboard-new-row",
"dashboard-new-request",
Some("user-1"),
Some("key-1"),
Some("Personal"),
"Provider",
"model-1",
100,
20,
1.0,
0.8,
now.timestamp(),
);
recent.request_metadata = Some(json!({"analytics_attribution":{"is_standalone":false}}));
let mut historical = recent.clone();
historical.id = "dashboard-old-row".into();
historical.request_id = "dashboard-old-request".into();
historical.created_at_unix_ms = (now - chrono::Duration::days(500)).timestamp() as u64;
let repository = Arc::new(
InMemoryUsageReadRepository::seed([recent, historical])
.with_dashboard_stats_since(since)
.with_analytics_users([
sample_user_summary("user-1", "Alice", "user", true),
sample_user_summary("user-2", "Bob", "user", false),
]),
);
let gateway = build_router_with_state(
AppState::new()
.unwrap()
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(repository)),
);
let (url, handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let endpoint = format!("{url}/api/admin/overview/dashboard/summary");
assert!(matches!(
client.get(&endpoint).send().await.unwrap().status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let response = admin_request(client.get(format!("{endpoint}?timezone=UTC")))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let value: serde_json::Value = response.json().await.unwrap();
assert_eq!(value["stats_since"], since.to_rfc3339());
assert_eq!(value["today"]["request_count"], 1);
assert_eq!(value["total"]["request_count"], 1);
assert_eq!(value["today"]["input_tokens"], 100);
assert_eq!(value["today"]["output_tokens"], 20);
assert_eq!(value["total"]["total_tokens"], 120);
assert_eq!(value["total"]["billable_amount"]["value"], "0.80000000");
assert_eq!(value["users"]["total"], 2);
assert_eq!(value["active_days"], 1);
assert_eq!(value["concurrency"]["scope"], "node");
assert!(value["today"].get("latency_p95_ms").is_none());
let invalid = admin_request(client.get(format!("{endpoint}?timezone=invalid")))
.send()
.await
.unwrap();
assert_eq!(invalid.status(), StatusCode::BAD_REQUEST);
handle.abort();
}
#[tokio::test]
async fn overview_dashboard_keeps_today_and_lifetime_totals_separate() {
let now = chrono::Utc::now();
let today_start = now.date_naive().and_hms_opt(0, 0, 0).unwrap().and_utc();
let row = sample_usage_row(
"today-row",
"today-request",
Some("user-1"),
Some("key-1"),
Some("Personal"),
"Provider",
"model-1",
100,
20,
1.0,
0.8,
now.timestamp()
.saturating_sub(1)
.max(today_start.timestamp()),
);
let mut historical = row.clone();
historical.id = "historical-row".into();
historical.request_id = "historical-request".into();
historical.created_at_unix_ms = (now - chrono::Duration::days(500)).timestamp() as u64;
let mut future = row.clone();
future.id = "future-row".into();
future.request_id = "future-request".into();
future.created_at_unix_ms = (now + chrono::Duration::days(1)).timestamp() as u64;
let repository = Arc::new(
InMemoryUsageReadRepository::seed([row, historical, future]).with_analytics_users([
sample_user_summary("user-1", "Alice", "user", true),
sample_user_summary("user-2", "Bob", "user", true),
]),
);
let gateway = build_router_with_state(
AppState::new()
.unwrap()
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(repository)),
);
let (url, handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let endpoint = format!("{url}/api/admin/overview/dashboard");
let anonymous = client.get(&endpoint).send().await.unwrap();
assert!(matches!(
anonymous.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let response = admin_request(client.get(format!("{endpoint}?timezone=UTC")))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(response.headers()[http::header::CACHE_CONTROL]
.to_str()
.unwrap()
.contains("no-store"));
let data: serde_json::Value = response.json().await.unwrap();
assert_eq!(data["today"]["data"]["request_count"], 1);
assert_eq!(data["total"]["data"]["request_count"], 2);
assert_eq!(data["today"]["data"]["total_tokens"], 120);
assert_eq!(data["total"]["data"]["total_tokens"], 240);
assert_eq!(
data["today"]["data"]["billable_amount"]["value"],
"0.80000000"
);
assert_eq!(
data["total"]["data"]["billable_amount"]["value"],
"1.60000000"
);
assert_eq!(data["today"]["data"]["enabled_users"], 2);
assert_eq!(data["today"]["meta"]["range"]["timezone"], "UTC");
assert_eq!(data["total"]["meta"]["range"]["period"], "all_time");
assert_eq!(
data["today"]["meta"]["read_revision"],
data["total"]["meta"]["read_revision"]
);
assert_eq!(
data["today"]["meta"]["range"]["to"],
data["total"]["meta"]["range"]["to"]
);
let total_endpoint = format!("{endpoint}/total");
let anonymous_total = client.get(&total_endpoint).send().await.unwrap();
assert!(matches!(
anonymous_total.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let pending = admin_request(client.get(format!("{total_endpoint}?timezone=UTC")))
.send()
.await
.unwrap();
assert_eq!(pending.status(), StatusCode::ACCEPTED);
assert_eq!(
pending.json::<serde_json::Value>().await.unwrap()["status"],
"pending"
);
let ready = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
let response =
admin_request(client.get(format!("{total_endpoint}?timezone=Asia%2FShanghai")))
.send()
.await
.unwrap();
if response.status() == StatusCode::OK {
break response.json::<serde_json::Value>().await.unwrap();
}
assert_eq!(response.status(), StatusCode::ACCEPTED);
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert_eq!(ready["status"], "ready");
assert_eq!(ready["stale"], false);
assert_eq!(ready["total"]["data"]["request_count"], 2);
assert_eq!(ready["total"]["data"]["total_tokens"], 240);
assert_eq!(
ready["total"]["data"]["billable_amount"]["value"],
"1.60000000"
);
assert_eq!(ready["total"]["meta"]["range"]["timezone"], "Asia/Shanghai");
let cached_utc = admin_request(client.get(format!("{total_endpoint}?timezone=UTC")))
.send()
.await
.unwrap()
.json::<serde_json::Value>()
.await
.unwrap();
assert_eq!(
cached_utc["total"]["meta"]["read_revision"],
ready["total"]["meta"]["read_revision"]
);
assert_eq!(
cached_utc["total"]["meta"]["generated_at"],
ready["total"]["meta"]["generated_at"]
);
for query in [
"from=2020-01-01T00:00:00Z",
"user_id=user-1",
"timezone=invalid",
"timezone=UTC&timezone=UTC",
] {
let response = admin_request(client.get(format!("{endpoint}?{query}")))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{query}");
}
handle.abort();
}
#[tokio::test]
async fn overview_preserves_scope_pagination_amounts_and_csv_across_precise_range() {
let from = chrono::DateTime::parse_from_rfc3339("2026-09-01T23:45:00Z")
.unwrap()
.timestamp();
let mut row = sample_usage_row(
"row-1",
"request-1",
Some("user-1"),
Some("key-1"),
Some("Personal"),
"Provider",
"model-1",
100,
20,
1.0,
0.8,
from + 60,
);
row.request_metadata =
Some(json!({"analytics_attribution": {"is_standalone": false, "record_kind": "request"}}));
let mut outside = row.clone();
outside.id = "row-outside".into();
outside.request_id = "request-outside".into();
outside.created_at_unix_ms = (from + 30 * 60) as u64;
let repository = Arc::new(
InMemoryUsageReadRepository::seed([row, outside]).with_analytics_users([
sample_user_summary("user-1", "Alice", "user", true),
sample_user_summary("user-2", "Bob", "user", true),
]),
);
let gateway = build_router_with_state(
AppState::new()
.unwrap()
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(repository)),
);
let (url, handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let anonymous = client
.get(format!("{url}/api/admin/overview/summary?{RANGE}"))
.send()
.await
.unwrap();
assert!(matches!(
anonymous.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let summary = admin_request(client.get(format!("{url}/api/admin/overview/summary?{RANGE}")))
.send()
.await
.unwrap();
assert_eq!(summary.status(), StatusCode::OK);
assert!(summary
.headers()
.get(http::header::CACHE_CONTROL)
.unwrap()
.to_str()
.unwrap()
.contains("no-store"));
let summary: serde_json::Value = summary.json().await.unwrap();
assert_eq!(summary["data"]["request_count"], 1);
assert_eq!(summary["data"]["billable_amount"]["value"], "0.80000000");
assert_eq!(summary["data"]["quota_covered_amount"]["status"], "unknown");
assert_eq!(summary["meta"]["range"]["time_basis"], "request_started_at");
let charts = admin_request(client.get(format!(
"{url}/api/admin/overview/dashboard/charts?{RANGE}&granularity=day"
)))
.send()
.await
.unwrap();
assert_eq!(charts.status(), StatusCode::OK);
let charts: serde_json::Value = charts.json().await.unwrap();
assert_eq!(charts["data"]["summary"]["request_count"], 1);
assert_eq!(charts["data"]["series"].as_array().unwrap().len(), 1);
assert_eq!(charts["data"]["models"][0]["id"], "model-1");
assert_eq!(
charts["data"]["models"][0]["billable_amount"]["value"],
"0.80000000"
);
assert_eq!(
charts["data"]["providers"][0]["billable_amount"]["value"],
"0.80000000"
);
let users = admin_request(client.get(format!(
"{url}/api/admin/overview/users?{RANGE}&limit=1&offset=1&sort=request_count"
)))
.send()
.await
.unwrap();
assert_eq!(users.status(), StatusCode::OK);
let users: serde_json::Value = users.json().await.unwrap();
assert_eq!(users["data"]["total"], 2);
assert_eq!(users["data"]["items"][0]["user_id"], "user-2");
assert_eq!(users["data"]["items"][0]["request_count"], 0);
assert_eq!(users["data"]["summary"]["user_count"], 2);
assert_eq!(users["data"]["summary"]["active_user_count"], 1);
assert_eq!(users["data"]["summary"]["request_count"], 1);
assert_eq!(
users["data"]["summary"]["billable_amount"]["value"],
"0.80000000"
);
assert!(users["data"]["finance_summary"].is_null());
assert!(users["data"]["items"][0]["finance"].is_null());
let csv = admin_request(client.get(format!(
"{url}/api/admin/overview/users?{RANGE}&format=csv&limit=1"
)))
.send()
.await
.unwrap();
assert_eq!(csv.status(), StatusCode::OK);
let csv = csv.text().await.unwrap();
assert!(csv.contains("Alice") && csv.contains("Bob"));
assert!(csv.contains("finance.recharge_amount.value"));
assert!(csv.contains("finance.plan_purchase_amount.value"));
let detail =
admin_request(client.get(format!("{url}/api/admin/overview/users/user-1?{RANGE}")))
.send()
.await
.unwrap();
assert_eq!(detail.status(), StatusCode::OK);
let detail: serde_json::Value = detail.json().await.unwrap();
assert_eq!(detail["meta"]["scope"]["kind"], "credential_owner");
assert_eq!(detail["data"]["summary"]["request_count"], 1);
assert!(detail["data"]["finance"].is_null());
assert!(detail["data"]["payments"].is_null());
let conflicting = admin_request(client.get(format!(
"{url}/api/admin/overview/users/user-1?{RANGE}&user_id=user-2&payment_limit=2"
)))
.send()
.await
.unwrap();
assert_eq!(conflicting.status(), StatusCode::BAD_REQUEST);
let consumption = admin_request(client.get(format!(
"{url}/api/admin/overview/consumption?{RANGE}&status=success&sort=started_at"
)))
.send()
.await
.unwrap();
assert_eq!(consumption.status(), StatusCode::OK);
let consumption: serde_json::Value = consumption.json().await.unwrap();
assert_eq!(consumption["data"]["items"][0]["id"], "row-1");
assert_eq!(consumption["data"]["items"][0]["request_id"], "request-1");
for (filter, expected) in [
(
"api_key_id=key-1&provider_id=provider-1&request_id=request-1&status=success",
1,
),
("api_key_id=other-key", 0),
("user_id=user-1&attribution_kind=employee", 1),
("user_id=user-2&attribution_kind=employee", 0),
("provider_id=other-provider", 0),
("slow_threshold_ms=5000", 0),
(
"slow_threshold_ms=350&is_stream=false&endpoint_kind=chat",
1,
),
] {
let records = admin_request(client.get(format!(
"{url}/api/admin/usage/records?{RANGE}&{filter}&include_total=true"
)))
.send()
.await
.unwrap();
assert_eq!(records.status(), StatusCode::OK, "filter: {filter}");
let records: serde_json::Value = records.json().await.unwrap();
assert_eq!(records["total"], expected, "filter: {filter}");
}
let invalid = admin_request(client.get(format!(
"{url}/api/admin/overview/summary?{RANGE}&made_up=1"
)))
.send()
.await
.unwrap();
assert_eq!(invalid.status(), StatusCode::BAD_REQUEST);
handle.abort();
}
@@ -1181,6 +1181,12 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
.as_deref()
.is_some_and(|value| value.starts_with("aether-auth-api-key-secret-v2:"))));
assert_eq!(recovery_payload["version"], "1.5");
assert!(recovery_payload["users"][0]["api_keys"][0]
.get("credential_kind")
.is_none());
assert!(recovery_payload["standalone_keys"][0]
.get("credential_kind")
.is_none());
assert_eq!(recovery_payload["users"][0]["password_hash"], "argon2-hash");
assert_eq!(
recovery_payload["users"][0]["api_keys"][0]["key_hash"],
@@ -1228,6 +1234,10 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["version"], "1.6");
assert!(payload["users"][0]["api_keys"][0]
.get("credential_kind")
.is_none());
assert!(payload["standalone_keys"][0].get("credential_kind").is_none());
assert!(payload["exported_at"].as_str().is_some());
assert_eq!(payload["user_groups"][0]["name"], "Restricted GPT");
assert!(payload["user_groups"][0].get("priority").is_none());
@@ -52,11 +52,193 @@ const TEST_EMAIL_VERIFICATION_TOKEN: &str =
#[path = "public_support/auth_cookie.rs"]
mod auth_cookie;
#[path = "public_support/announcement_user_list.rs"]
mod announcement_user_list;
#[path = "public_support/dashboard.rs"]
mod dashboard;
#[path = "public_support/vscodex.rs"]
mod vscodex;
#[tokio::test]
async fn health_v2_public_scope_filters_summary_lists_and_details_before_projection() {
let now = Utc::now();
let mut published = sample_user_usage_audit(
"private-detail-id",
"private-request-id",
"private-user-id",
"internal-published-model",
"private-provider-name",
"failed",
now - chrono::Duration::minutes(5),
);
published.request_metadata = Some(json!({
"analytics_failure": {"origin": "upstream", "stage": "response", "reason": "private-diagnostic", "schema_version": 1},
}));
let hidden = sample_user_usage_audit(
"private-hidden-id",
"private-hidden-request",
"private-user-id",
"unpublished-model",
"private-provider-name",
"completed",
now - chrono::Duration::minutes(4),
);
let data = GatewayDataState::with_usage_reader_for_tests(Arc::new(InMemoryUsageReadRepository::seed(vec![published, hidden])))
.with_system_config_values_for_tests(vec![("health_publication_v1".into(), json!({
"enabled": true, "objects": [
{"public_id": "model-api", "kind": "model", "value": "internal-published-model", "display_name": "Model API"},
],
}))]);
let (gateway_url, gateway_handle) = start_server(build_router_with_state(
AppState::new().unwrap().with_data_state_for_tests(data),
))
.await;
let client = reqwest::Client::new();
for resource in ["summary", "objects", "objects/model-api"] {
let response = client
.get(format!(
"{gateway_url}/api/public/health/v2/{resource}?kind=model"
))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body: serde_json::Value = response.json().await.unwrap();
let metrics = match resource {
"summary" => &body["data"]["requests"],
"objects" => &body["data"]["items"][0],
_ => &body["data"],
};
assert_eq!(
metrics["request_count"], 1,
"{resource} includes only the published model"
);
if resource == "objects" {
assert_eq!(body["data"]["total"], 1);
assert!(!metrics["timeline"].as_array().unwrap().is_empty());
}
let encoded = body.to_string();
for forbidden in [
"private-",
"internal-published-model",
"unpublished-model",
"provider_id",
"api_key_id",
"source_value",
"attempts",
"error_message",
"analytics_failure",
] {
assert!(
!encoded.contains(forbidden),
"{resource} leaked {forbidden}"
);
}
}
let empty_scope = client
.get(format!(
"{gateway_url}/api/public/health/v2/summary?kind=api_format"
))
.send()
.await
.unwrap();
assert_eq!(empty_scope.status(), StatusCode::OK);
let empty: serde_json::Value = empty_scope.json().await.unwrap();
assert_eq!(empty["data"]["object_count"], 0);
assert_eq!(empty["data"]["requests"]["request_count"], 0);
gateway_handle.abort();
}
#[tokio::test]
async fn health_v2_authenticated_user_access_does_not_publish_anonymous_status() {
let now = Utc::now();
let user = sample_auth_user(now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".into(), json!(user.id)),
("role".into(), json!(user.role)),
(
"created_at".into(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".into(), json!("health-user-session")),
]),
now + chrono::Duration::hours(1),
);
let data = GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed(vec![sample_user_usage_audit(
"health-user-row",
"health-user-request",
&user.id,
"model-health",
"internal-provider",
"completed",
now - chrono::Duration::minutes(5),
)]),
))
.with_user_reader(Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
user,
])))
.with_provider_catalog_reader(Arc::new(InMemoryProviderCatalogReadRepository::seed(
Vec::new(),
Vec::new(),
Vec::new(),
)))
.with_system_config_values_for_tests(vec![(
"health_publication_v1".into(),
json!({ "enabled": false, "objects": [] }),
)]);
let state = AppState::new()
.unwrap()
.with_data_state_for_tests(data)
.with_auth_session_for_tests(sample_auth_session(
"user-auth-1",
"health-user-session",
"health-user-device",
"refresh-token-placeholder",
now,
));
let (gateway_url, gateway_handle) = start_server(build_router_with_state(state)).await;
let client = reqwest::Client::new();
let url = format!("{gateway_url}/api/users/me/health/v2/objects?kind=api_format");
let anonymous = client.get(&url).send().await.unwrap();
assert_eq!(anonymous.status(), StatusCode::UNAUTHORIZED);
let authenticated = client
.get(&url)
.bearer_auth(&access_token)
.header("x-client-device-id", "health-user-device")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.unwrap();
assert_eq!(authenticated.status(), StatusCode::OK);
let body: serde_json::Value = authenticated.json().await.unwrap();
assert_eq!(body["meta"]["scope"]["kind"], "authenticated");
assert_eq!(body["data"]["total"], 1);
assert_eq!(body["data"]["items"][0]["request_count"], 1);
assert!(body["data"]["items"][0].get("attempts").is_none());
assert!(!body.to_string().contains("internal-provider"));
let provider = client
.get(format!(
"{gateway_url}/api/users/me/health/v2/objects?kind=provider"
))
.bearer_auth(&access_token)
.header("x-client-device-id", "health-user-device")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.unwrap();
assert_eq!(provider.status(), StatusCode::BAD_REQUEST);
let public = client
.get(format!("{gateway_url}/api/public/health/v2/objects"))
.send()
.await
.unwrap();
assert_eq!(public.status(), StatusCode::NOT_FOUND);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_announcements_list_without_proxying_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -7806,6 +7988,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
.expect("created id should be string")
.to_string();
assert_eq!(create_payload["name"], "writer-key");
assert!(create_payload.get("credential_kind").is_none());
assert_eq!(create_payload["rate_limit"], 120);
assert_eq!(create_payload["concurrent_limit"], serde_json::Value::Null);
assert_eq!(
@@ -7847,6 +8030,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
.await
.expect("json body should parse");
assert_eq!(update_payload["name"], "writer-key-renamed");
assert!(update_payload.get("credential_kind").is_none());
assert_eq!(update_payload["rate_limit"], 30);
assert_eq!(update_payload["concurrent_limit"], 4);
assert_eq!(
@@ -0,0 +1,317 @@
use super::*;
struct AnnouncementUserFixture {
url: String,
token: String,
device: String,
client: reqwest::Client,
upstream_hits: Arc<Mutex<usize>>,
gateway: tokio::task::JoinHandle<()>,
upstream: tokio::task::JoinHandle<()>,
}
impl AnnouncementUserFixture {
async fn start(
user_id: &str,
role: &str,
repository: Arc<InMemoryAnnouncementReadRepository>,
) -> Self {
let now = Utc::now();
let mut user = sample_auth_user(now);
user.id = user_id.into();
user.role = role.into();
let session = format!("{user_id}-session");
let device = format!("{user_id}-device");
let token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".into(), json!(user.id)),
("role".into(), json!(user.role)),
(
"created_at".into(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".into(), json!(session)),
]),
now + chrono::Duration::hours(1),
);
let (url, upstream_hits, gateway, upstream) = start_auth_announcement_gateway_with_state(
user,
sample_auth_wallet(user_id, now),
[sample_auth_session(
user_id,
&session,
&device,
"refresh-placeholder",
now,
)],
repository,
)
.await;
Self {
url,
token,
device,
client: reqwest::Client::new(),
upstream_hits,
gateway,
upstream,
}
}
fn auth(&self, request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
request
.bearer_auth(&self.token)
.header("x-client-device-id", &self.device)
.header("user-agent", "AetherTest/1.0")
}
async fn list(&self, suffix: &str) -> serde_json::Value {
let response = self
.auth(
self.client
.get(format!("{}/api/announcements/users/me{suffix}", self.url)),
)
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK, "{suffix}");
response.json().await.unwrap()
}
async fn mark_read(&self, id: &str) {
let response = self
.auth(
self.client
.patch(format!("{}/api/announcements/{id}/read-status", self.url)),
)
.json(&json!({ "is_read": true }))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
async fn read_all(&self) {
let response = self
.auth(
self.client
.post(format!("{}/api/announcements/read-all", self.url)),
)
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
}
impl Drop for AnnouncementUserFixture {
fn drop(&mut self) {
self.gateway.abort();
self.upstream.abort();
}
}
fn announcement(id: &str, pinned: bool, priority: i32, created_at: i64) -> StoredAnnouncement {
StoredAnnouncement::new(
id.into(),
format!("Title {id}"),
format!("Content {id}"),
"info".into(),
priority,
true,
pinned,
false,
Some("author-1".into()),
Some("Author".into()),
None,
None,
created_at,
created_at,
)
.unwrap()
}
fn ids(payload: &serde_json::Value) -> Vec<&str> {
payload["items"]
.as_array()
.unwrap()
.iter()
.map(|item| item["id"].as_str().unwrap())
.collect()
}
#[tokio::test]
async fn announcement_user_list_keeps_visible_order_and_global_unread_counts_across_pages() {
let now = Utc::now().timestamp();
let mut draft = announcement("draft", true, 999, now);
draft.is_active = false;
let mut future = announcement("future", true, 999, now);
future.start_time_unix_secs = Some((now + 3600) as u64);
let mut expired = announcement("expired", true, 999, now);
expired.end_time_unix_secs = Some((now - 3600) as u64);
let repository = Arc::new(InMemoryAnnouncementReadRepository::seed_with_reads(
vec![
announcement("active-e", false, 100, now - 20),
announcement("active-b", true, 10, now - 100),
draft,
announcement("active-d", false, 100, now - 10),
future,
announcement("active-c", true, 5, now - 50),
expired,
announcement("active-a", true, 10, now - 100),
],
[("list-user".into(), "active-b".into())],
));
let fixture = AnnouncementUserFixture::start("list-user", "user", repository).await;
let all = fixture.list("").await;
assert_eq!(
ids(&all),
["active-a", "active-b", "active-c", "active-d", "active-e"]
);
assert_eq!(all["total"], 5);
assert_eq!(all["unread_count"], 4);
assert_eq!(all["limit"], 20);
assert_eq!(all["offset"], 0);
assert_eq!(all["items"][0]["is_read"], false);
assert_eq!(all["items"][1]["is_read"], true);
assert_eq!(all["items"][0]["content"], "Content active-a");
assert_eq!(all["items"][0]["author"]["username"], "Author");
let page = fixture.list("?limit=2&offset=1&unread_only=false").await;
assert_eq!(ids(&page), ["active-b", "active-c"]);
assert_eq!(page["total"], 5);
assert_eq!(page["unread_count"], 4);
assert_eq!(page["limit"], 2);
assert_eq!(page["offset"], 1);
let unread = fixture.list("?limit=2&offset=1&unread_only=true").await;
assert_eq!(ids(&unread), ["active-c", "active-d"]);
assert_eq!(unread["total"], 4);
assert_eq!(unread["unread_count"], 4);
for (query, total) in [
("?limit=2&offset=999", 5),
("?limit=2&offset=999&unread_only=true", 4),
] {
let outside = fixture.list(query).await;
assert!(ids(&outside).is_empty());
assert_eq!(outside["total"], total);
assert_eq!(outside["unread_count"], 4);
}
fixture.mark_read("active-d").await;
let unread = fixture.list("?unread_only=true").await;
assert_eq!(ids(&unread), ["active-a", "active-c", "active-e"]);
assert_eq!(unread["total"], 3);
assert_eq!(unread["unread_count"], 3);
let all = fixture.list("").await;
assert_eq!(all["items"][3]["is_read"], true);
let badge = fixture
.auth(fixture.client.get(format!(
"{}/api/announcements/users/me/unread-count",
fixture.url
)))
.send()
.await
.unwrap();
assert_eq!(badge.status(), StatusCode::OK);
assert_eq!(
badge.json::<serde_json::Value>().await.unwrap()["unread_count"],
3
);
fixture.read_all().await;
let all = fixture.list("").await;
assert_eq!(all["total"], 5);
assert_eq!(all["unread_count"], 0);
assert!(all["items"]
.as_array()
.unwrap()
.iter()
.all(|item| item["is_read"] == true));
let unread = fixture.list("?unread_only=true").await;
assert!(ids(&unread).is_empty());
assert_eq!(unread["total"], 0);
assert_eq!(*fixture.upstream_hits.lock().unwrap(), 0);
}
#[tokio::test]
async fn announcement_user_list_is_personal_for_user_admin_and_audit_admin() {
let now = Utc::now().timestamp();
let repository = Arc::new(InMemoryAnnouncementReadRepository::seed(vec![
announcement("shared-notice", false, 1, now),
announcement("second-notice", false, 0, now),
]));
for role in ["user", "admin", "audit_admin"] {
let fixture = AnnouncementUserFixture::start(
&format!("announcement-{role}"),
role,
Arc::clone(&repository),
)
.await;
let initial = fixture.list("").await;
assert_eq!(initial["total"], 2, "{role}");
assert_eq!(
initial["unread_count"], 2,
"another user's reads must not affect {role}"
);
assert!(initial["items"]
.as_array()
.unwrap()
.iter()
.all(|item| item["is_read"] == false));
fixture.mark_read("shared-notice").await;
let changed = fixture.list("").await;
assert_eq!(changed["unread_count"], 1, "{role}");
assert_eq!(changed["items"][0]["is_read"], true);
fixture.read_all().await;
assert_eq!(fixture.list("").await["unread_count"], 0, "{role}");
assert_eq!(*fixture.upstream_hits.lock().unwrap(), 0);
}
}
#[tokio::test]
async fn announcement_user_list_requires_auth_and_rejects_invalid_filters() {
let fixture = AnnouncementUserFixture::start(
"announcement-validation",
"user",
Arc::new(InMemoryAnnouncementReadRepository::seed(Vec::new())),
)
.await;
let url = format!("{}/api/announcements/users/me", fixture.url);
assert_eq!(
fixture.client.get(&url).send().await.unwrap().status(),
StatusCode::UNAUTHORIZED
);
for query in [
"limit=0",
"limit=101",
"limit=-1",
"limit=1.5",
"limit=",
"offset=-1",
"offset=1.5",
"offset=9223372036854775808",
"offset=18446744073709551616",
"unread_only=invalid",
"unread_only=",
"limit=20&limit=20",
"offset=0&offset=0",
"unread_only=false&unread_only=false",
"active_only=false",
"user_id=another-user",
"now=4102444800",
] {
let response = fixture
.auth(fixture.client.get(format!("{url}?{query}")))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{query}");
}
let empty = fixture.list("/?limit=100&offset=9223372036854775807").await;
assert!(ids(&empty).is_empty());
assert_eq!(empty["total"], 0);
assert_eq!(empty["unread_count"], 0);
assert_eq!(empty["limit"], 100);
assert_eq!(empty["offset"], i64::MAX);
assert_eq!(*fixture.upstream_hits.lock().unwrap(), 0);
}
@@ -39,6 +39,7 @@ const INTERNAL_REPORT_OBSERVATION_FIELDS: &[&str] = &[
"client_response_headers",
"upstream_response",
"error_flow",
"analytics_failure",
"transport_error",
"input_tokens",
"cache_creation_input_tokens",
@@ -0,0 +1,291 @@
use aether_usage_runtime::{
stream_report_missing_terminal_event, stream_report_represents_failure,
sync_report_represents_failure, GatewayStreamReportRequest, GatewaySyncReportRequest,
};
use serde_json::{json, Map, Value};
pub(crate) fn execution_error_analytics_context(
context: Option<&Value>,
error: &aether_contracts::ExecutionError,
) -> Option<Value> {
use aether_contracts::{ExecutionErrorKind as Kind, ExecutionPhase as Phase};
let stage = match error.phase {
Phase::Connect => "connect",
Phase::Handshake => "handshake",
Phase::Write => "request_write",
Phase::FirstByte => "first_byte",
Phase::StreamRead => "stream_read",
Phase::Decode => "decode",
Phase::Finalize => "finalize",
};
let (origin, reason) = match error.kind {
Kind::ConnectTimeout => ("transport", "connect_timeout"),
Kind::FirstByteTimeout => ("upstream", "first_byte_timeout"),
Kind::ReadTimeout => ("transport", "read_timeout"),
Kind::TlsError => ("transport", "tls_error"),
Kind::ProxyError => ("transport", "proxy_error"),
Kind::Upstream4xx | Kind::Upstream5xx => ("upstream", "upstream_response_error"),
Kind::ProtocolError => ("upstream", "protocol_error"),
Kind::Internal => ("gateway", "execution_internal_error"),
Kind::Cancelled => ("unknown", "execution_cancelled"),
};
with_analytics_failure(context, origin, stage, reason)
}
pub(crate) fn with_analytics_failure(
context: Option<&Value>,
origin: &'static str,
stage: &'static str,
reason: &'static str,
) -> Option<Value> {
let mut object = context
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
object.insert(
"analytics_failure".into(),
json!({
"origin": origin, "stage": stage, "reason": reason, "schema_version": 1,
}),
);
Some(Value::Object(object))
}
fn normalized_failure_context(context: Option<&Value>, failed: bool) -> Map<String, Value> {
let mut object = context
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
if !failed {
// A successful retry supersedes a previous candidate's failure classification.
object.remove("analytics_failure");
object.remove("error_flow");
object.remove("transport_error");
}
object
}
fn classify_observed_failure(object: &Map<String, Value>, stage: &'static str) -> Option<Value> {
if object
.get("analytics_failure")
.and_then(Value::as_object)
.is_some()
{
return Some(Value::Object(object.clone()));
}
if object.get("transport_error").and_then(Value::as_bool) == Some(true) {
return with_analytics_failure(
Some(&Value::Object(object.clone())),
"transport",
stage,
"upstream_transport_error",
);
}
if object.get("error_flow").is_some_and(|flow| {
flow.get("source").and_then(Value::as_str) == Some("upstream_response")
&& flow
.get("status_code")
.and_then(Value::as_u64)
.is_some_and(|status| status >= 400)
}) || object
.get("upstream_response")
.and_then(|value| value.get("status_code"))
.and_then(Value::as_u64)
.is_some_and(|status| status >= 400)
{
return with_analytics_failure(
Some(&Value::Object(object.clone())),
"upstream",
stage,
"upstream_response_error",
);
}
None
}
pub(crate) fn sync_analytics_context(
context: Option<&Value>,
payload: &GatewaySyncReportRequest,
) -> Option<Value> {
let failed = sync_report_represents_failure(payload, None);
let object = normalized_failure_context(context, failed);
if !failed {
return Some(Value::Object(object));
}
classify_observed_failure(&object, "response").or(Some(Value::Object(object)))
}
pub(crate) fn stream_analytics_context(
context: Option<&Value>,
payload: &GatewayStreamReportRequest,
downstream_cancelled: bool,
) -> Option<Value> {
if downstream_cancelled {
return with_analytics_failure(context, "client", "delivery", "downstream_disconnect");
}
let failed = stream_report_represents_failure(payload);
let object = normalized_failure_context(context, failed);
if !failed {
return Some(Value::Object(object));
}
if let Some(classified) = classify_observed_failure(&object, "stream_read") {
return Some(classified);
}
if payload
.terminal_summary
.as_ref()
.is_some_and(|summary| summary.parser_error.is_some())
{
return with_analytics_failure(
Some(&Value::Object(object)),
"gateway",
"decode",
"response_decode_error",
);
}
if stream_report_missing_terminal_event(payload) {
return with_analytics_failure(
Some(&Value::Object(object)),
"upstream",
"stream_read",
"missing_terminal_event",
);
}
Some(Value::Object(object))
}
pub(crate) fn gateway_error_analytics_context(
context: Option<&Value>,
error: &crate::GatewayError,
) -> Option<Value> {
match error {
crate::GatewayError::AdmissionTimeout { .. } => {
with_analytics_failure(context, "gateway", "admission", "gateway_admission_timeout")
}
crate::GatewayError::LocalExecutionPlanningTimeout { .. } => {
with_analytics_failure(context, "gateway", "routing", "planning_timeout")
}
crate::GatewayError::PlanUsageLimited(_) => {
with_analytics_failure(context, "client", "admission", "quota_exceeded")
}
crate::GatewayError::UpstreamUnavailable { .. } => {
with_analytics_failure(context, "upstream", "connect", "upstream_unavailable")
}
crate::GatewayError::ControlUnavailable { .. } => {
with_analytics_failure(context, "gateway", "routing", "control_unavailable")
}
crate::GatewayError::Internal(_) => {
with_analytics_failure(context, "gateway", "finalize", "internal_error")
}
_ => context.cloned(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::BTreeMap;
fn sync_payload(status: u16) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id: "t".into(),
report_kind: "openai_chat_sync_success".into(),
report_context: None,
status_code: status,
headers: BTreeMap::new(),
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
}
}
#[test]
fn analytics_failure_preserves_upstream_auth_and_throttling_as_service_failures() {
for status in [401, 429] {
let context =
json!({"error_flow": {"source": "upstream_response", "status_code": status}});
let result = sync_analytics_context(Some(&context), &sync_payload(status)).unwrap();
assert_eq!(result["analytics_failure"]["origin"], "upstream");
assert_eq!(result["analytics_failure"]["schema_version"], 1);
let unknown = sync_analytics_context(Some(&json!({})), &sync_payload(status)).unwrap();
assert!(unknown.get("analytics_failure").is_none());
}
}
#[test]
fn analytics_failure_successful_retry_clears_previous_observations() {
let context = json!({"request_id": "keep", "analytics_failure": {"origin": "upstream"}, "error_flow": {"source": "upstream_response"}, "transport_error": true});
let result = sync_analytics_context(Some(&context), &sync_payload(200)).unwrap();
assert_eq!(result["request_id"], "keep");
assert!(result.get("analytics_failure").is_none());
assert!(result.get("transport_error").is_none());
}
#[test]
fn analytics_failure_transport_uses_observation_and_not_public_error_text() {
let context =
json!({"transport_error": true, "error_flow": {"source": "upstream_response"}});
let result = sync_analytics_context(Some(&context), &sync_payload(502)).unwrap();
assert_eq!(result["analytics_failure"]["origin"], "transport");
assert_eq!(
result["analytics_failure"]["reason"],
"upstream_transport_error"
);
}
#[test]
fn analytics_failure_gateway_admission_is_not_client_quota() {
let error = crate::GatewayError::AdmissionTimeout {
trace_id: "t".into(),
gate: "upstream",
queue_budget_ms: 20,
};
let result = gateway_error_analytics_context(None, &error).unwrap();
assert_eq!(result["analytics_failure"]["origin"], "gateway");
assert_eq!(result["analytics_failure"]["stage"], "admission");
}
#[test]
fn analytics_failure_structured_error_has_priority_over_http_diagnostic() {
let context = json!({"error_flow": {"source": "upstream_response", "status_code": 502}});
let error = aether_contracts::ExecutionError {
kind: aether_contracts::ExecutionErrorKind::Internal,
phase: aether_contracts::ExecutionPhase::Decode,
message: "secret diagnostic".into(),
upstream_status: None,
retryable: false,
failover_recommended: false,
};
let context = execution_error_analytics_context(Some(&context), &error).unwrap();
let result = sync_analytics_context(Some(&context), &sync_payload(502)).unwrap();
assert_eq!(result["analytics_failure"]["origin"], "gateway");
assert_eq!(result["analytics_failure"]["stage"], "decode");
assert!(!result.to_string().contains("secret diagnostic"));
}
#[test]
fn analytics_failure_downstream_disconnect_is_distinct_from_unclassified_cancellation() {
let payload = GatewayStreamReportRequest {
trace_id: "t".into(),
report_kind: "openai_chat_stream_error".into(),
report_context: None,
status_code: 499,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
};
let result = stream_analytics_context(None, &payload, true).unwrap();
assert_eq!(result["analytics_failure"]["origin"], "client");
assert_eq!(
result["analytics_failure"]["reason"],
"downstream_disconnect"
);
let unknown = sync_analytics_context(None, &sync_payload(499)).unwrap();
assert!(unknown.get("analytics_failure").is_none());
}
}
@@ -13,6 +13,7 @@ use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_USAGE_SYNC_REPORT};
use crate::{AppState, GatewayError};
mod context;
pub(crate) mod failure;
pub(crate) use context::{
attach_internal_gateway_report_capability, resolve_bound_internal_gateway_report_context,
};
@@ -677,6 +678,7 @@ mod tests {
),
("upstream_response".to_string(), json!({"id": "resp-123"})),
("error_flow".to_string(), json!({"stage": "upstream"})),
("analytics_failure".to_string(), json!({"origin": "upstream", "stage": "response", "reason": "upstream_response_error", "schema_version": 1})),
(
"client_response_headers".to_string(),
json!({"content-type": "application/json"}),
@@ -860,6 +862,27 @@ mod tests {
assert!(resolved.is_none(), "cross-operation use must be rejected");
}
for (field, value) in [
(
"analytics_attribution",
json!({"is_standalone": false, "actor_user_id": "forged-employee"}),
),
("analytics_measurement", json!({"source": "reported"})),
("usage_token_source", json!("estimated")),
] {
let mut forged = minted.clone();
forged[field] = value;
let resolved = resolve_bound_internal_gateway_report_context(
&state,
"trace-capability-fields-123",
"openai_video_create_sync_finalize",
Some(&forged),
)
.await
.expect("capability lookup should succeed");
assert!(resolved.is_none(), "unbound {field} must be rejected");
}
let valid = resolve_bound_internal_gateway_report_context(
&state,
"trace-capability-fields-123",