Merge main into Claude Code dynamic client profile

This commit is contained in:
dalamudx
2026-10-05 13:51:07 +08:00
402 changed files with 37328 additions and 10782 deletions
@@ -19,12 +19,12 @@ pub(crate) use crate::ai_serving::{
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
OPENAI_MEMORIES_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily;
@@ -6,12 +6,13 @@ use crate::ai_serving::planner::common::{
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_MEMORIES_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
@@ -101,7 +102,9 @@ fn build_sync_plan_payload_from_decision(
OPENAI_RESPONSES_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, false)?
}
OPENAI_IMAGE_SYNC_PLAN_KIND | OPENAI_SEARCH_SYNC_PLAN_KIND => {
OPENAI_IMAGE_SYNC_PLAN_KIND
| OPENAI_SEARCH_SYNC_PLAN_KIND
| OPENAI_MEMORIES_SYNC_PLAN_KIND => {
build_passthrough_sync_plan_from_decision(parts, payload)?
}
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
@@ -17,7 +17,7 @@ use crate::ai_serving::transport::ProviderOutboundRequestContext;
use crate::ai_serving::{
ClientSurface, ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot,
GatewayCredentialCarrier, GatewayProviderTransportSnapshot, PlannerAppState,
CODEX_RESPONSES_LITE_HEADER,
CODEX_RESPONSES_LITE_HEADER, OPENAI_MEMORIES_SYNC_PLAN_KIND,
};
use crate::cache::CacheLoadObserver;
use crate::client_session_affinity::client_session_affinity_from_api_request;
@@ -123,6 +123,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
transport: Option<&GatewayProviderTransportSnapshot>,
websocket_continuation: bool,
) -> Result<(), GatewayError> {
let native_memories = decision.decision_kind.as_deref() == Some(OPENAI_MEMORIES_SYNC_PLAN_KIND);
let provider_api_format = decision
.provider_api_format
.clone()
@@ -150,7 +151,12 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
input.requested_model.as_str(),
)
});
crate::ai_serving::apply_codex_openai_responses_lite_header_for_request_body_with_capabilities(
if native_memories {
decision
.provider_request_headers
.retain(|name, _| !name.eq_ignore_ascii_case(CODEX_RESPONSES_LITE_HEADER));
} else {
crate::ai_serving::apply_codex_openai_responses_lite_header_for_request_body_with_capabilities(
&mut decision.provider_request_headers,
decision.provider_request_body.as_ref(),
provider_type.as_str(),
@@ -159,6 +165,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
input.requested_model.as_str(),
model_capabilities.as_ref(),
);
}
let Some(context) = input.routing_context.as_ref() else {
// Cache identity headers are projected only at the terminal boundary. Any non-empty
@@ -260,7 +267,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
provider_headers.insert(HeaderName::from_static(name), value);
}
}
if original_provider_request_body.is_some() {
if original_provider_request_body.is_some() && !native_memories {
let provider_model = provider_request_body
.get("model")
.and_then(Value::as_str)
@@ -318,6 +325,12 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
}
.map_err(|_| invalid_routing_provider_contract())?;
}
if native_memories {
crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy(
&mut provider_request_body,
Some(crate::ai_serving::ApiOperation::OpenAiMemoriesSummarize),
);
}
let provider_model = provider_request_body
.get("model")
.and_then(Value::as_str)
@@ -339,7 +352,11 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
provider_type.as_str(),
provider_api_format.as_str(),
);
crate::ai_serving::apply_codex_openai_responses_lite_header_for_request_body_with_capabilities(
if native_memories {
provider_request_headers
.retain(|name, _| !name.eq_ignore_ascii_case(CODEX_RESPONSES_LITE_HEADER));
} else {
crate::ai_serving::apply_codex_openai_responses_lite_header_for_request_body_with_capabilities(
&mut provider_request_headers,
Some(&provider_request_body),
provider_type.as_str(),
@@ -348,6 +365,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
input.requested_model.as_str(),
model_capabilities.as_ref(),
);
}
crate::ai_serving::apply_codex_openai_compact_terminal_headers(
&mut provider_request_headers,
provider_type.as_str(),
@@ -382,6 +400,15 @@ fn apply_provider_outbound_request_policies_to_decision(
let Some(context) = input.provider_outbound_context.as_ref() else {
return;
};
let native_context;
let context = if decision.decision_kind.as_deref() == Some(OPENAI_MEMORIES_SYNC_PLAN_KIND) {
native_context = context
.clone()
.with_api_operation(crate::ai_serving::ApiOperation::OpenAiMemoriesSummarize);
&native_context
} else {
context
};
let results = crate::ai_serving::transport::apply_provider_outbound_request_policies(
transport,
provider_api_format,
@@ -24,7 +24,9 @@ use crate::ai_serving::transport::{
SameFormatProviderCompatibilityEditAction, SameFormatProviderHeadersInput,
GEMINI_CLI_USER_AGENT, GROK_CHAT_PATH,
};
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::ai_serving::{
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, CODEX_RESPONSES_LITE_HEADER,
};
use crate::{AppState, GatewayError};
mod policy;
@@ -255,7 +257,9 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
// re-enforce stream-field policy afterward.
// Kiro behavior classification already hard-requires upstream streaming,
// and the Kiro envelope does not use a top-level body stream field.
if prepared.kiro_auth.is_none() {
if prepared.kiro_auth.is_none()
&& spec.operation != Some(crate::ai_serving::ApiOperation::OpenAiMemoriesSummarize)
{
enforce_provider_body_stream_policy(
&mut base_provider_request_body,
prepared.provider_api_format.as_str(),
@@ -275,7 +279,8 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
prepared.mapped_model.as_str(),
source_model,
);
if let Err(violation) =
if spec.operation != Some(crate::ai_serving::ApiOperation::OpenAiMemoriesSummarize) {
if let Err(violation) =
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_replay_policy(
&mut base_provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
@@ -313,6 +318,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.await;
return Ok(None);
}
}
// Same-format requests skip `apply_transport_request_body_semantics`, so the opt-in
// Claude Code body mimicry has to be applied here as well.
@@ -597,6 +603,11 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
source_model,
codex_model_capabilities.as_ref(),
);
if spec.operation == Some(crate::ai_serving::ApiOperation::OpenAiMemoriesSummarize) {
provider_request_headers
.retain(|name, _| !name.eq_ignore_ascii_case(CODEX_RESPONSES_LITE_HEADER));
provider_request_headers.insert("accept".to_string(), "application/json".to_string());
}
crate::ai_serving::transport::xai::insert_cli_identity_headers_if_needed(
transport.as_ref(),
prepared.provider_api_format.as_str(),
@@ -3,7 +3,7 @@ use serde_json::Value;
use super::super::LocalSameFormatProviderSpec;
use crate::ai_serving::transport::{
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
build_same_format_provider_request_body_with_compatibility_report_and_reasoning_replay_policy as build_same_format_provider_request_body_with_compatibility_report_impl,
build_same_format_provider_request_body_for_operation as build_same_format_provider_request_body_with_compatibility_report_impl,
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
SameFormatProviderRequestBodyOutput,
};
@@ -69,6 +69,7 @@ pub(crate) fn build_same_format_provider_request_body_with_compatibility_report(
enable_model_directives,
},
reasoning_replay_policy,
spec.operation,
)
}
@@ -112,6 +112,22 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
Some(provider_request_body)
}
pub(crate) fn build_cross_format_openai_chat_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<String> {
crate::ai_serving::transport::build_cross_format_openai_chat_upstream_url(
transport,
mapped_model,
provider_api_format,
upstream_is_stream,
parts.uri.query(),
)
}
#[cfg(test)]
mod antigravity_schema_tests {
use super::*;
@@ -147,19 +163,3 @@ mod antigravity_schema_tests {
}
}
}
pub(crate) fn build_cross_format_openai_chat_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<String> {
crate::ai_serving::transport::build_cross_format_openai_chat_upstream_url(
transport,
mapped_model,
provider_api_format,
upstream_is_stream,
parts.uri.query(),
)
}
@@ -275,6 +275,24 @@ pub(crate) fn build_local_openai_responses_upstream_url(
)
}
pub(crate) fn build_cross_format_openai_responses_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<String> {
crate::ai_serving::transport::build_cross_format_openai_responses_upstream_url(
transport,
mapped_model,
client_api_format,
provider_api_format,
upstream_is_stream,
parts.uri.query(),
)
}
#[cfg(test)]
mod antigravity_schema_tests {
use super::*;
@@ -309,21 +327,3 @@ mod antigravity_schema_tests {
}
}
}
pub(crate) fn build_cross_format_openai_responses_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<String> {
crate::ai_serving::transport::build_cross_format_openai_responses_upstream_url(
transport,
mapped_model,
client_api_format,
provider_api_format,
upstream_is_stream,
parts.uri.query(),
)
}
@@ -163,7 +163,7 @@ pub(crate) use aether_ai_formats::api::{
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_MEMORIES_SYNC_PLAN_KIND, OPENAI_REALTIME_STREAM_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
@@ -78,7 +78,7 @@ pub(crate) use aether_provider_transport::{
build_local_openai_chat_upstream_url, build_local_openai_responses_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url, build_passthrough_headers,
build_request_trace_proxy_value, build_same_format_provider_headers,
build_same_format_provider_request_body,
build_same_format_provider_request_body, build_same_format_provider_request_body_for_operation,
build_same_format_provider_request_body_with_compatibility_report,
build_same_format_provider_request_body_with_compatibility_report_and_reasoning_replay_policy,
build_same_format_provider_upstream_url, build_standard_plan_fallback_headers,
+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());
}
}
+1
View File
@@ -126,6 +126,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens",
"/v1/responses",
"/v1/responses/compact",
"/v1/memories/trace_summarize",
"/v1/realtime",
"/v1/realtime/calls",
"/v1/live",
@@ -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,
@@ -727,3 +736,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,
}
}
@@ -88,6 +88,11 @@ pub(super) fn classify_ai_public_route(
true,
))
}
} else if method == http::Method::POST && normalized_path == "/v1/memories/trace_summarize" {
Some(
classified("ai_public", "openai", "memories", "openai:responses", true)
.with_api_operation(ApiOperation::OpenAiMemoriesSummarize),
)
} else if method == http::Method::POST && normalized_path == "/v1/alpha/search" {
Some(classified(
"ai_public",
@@ -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,
@@ -43,6 +43,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,
@@ -366,6 +369,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,
@@ -386,6 +409,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,
@@ -1729,6 +1764,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,
@@ -2817,6 +2906,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,
@@ -2889,6 +3007,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_key(|(provider, _)| *provider);
let mut models: Vec<_> = self.models.iter().collect();
models.sort_unstable_by_key(|(model, _)| *model);
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);
@@ -191,6 +191,7 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind(
let report_kind = match plan_kind {
"openai_chat_sync" => "openai_chat_sync_finalize",
"openai_responses_sync" => "openai_responses_sync_finalize",
"openai_memories_sync" => "openai_memories_sync_finalize",
"openai_responses_compact_sync" => "openai_responses_compact_sync_finalize",
"claude_chat_sync" => "claude_chat_sync_finalize",
"gemini_chat_sync" => "gemini_chat_sync_finalize",
@@ -576,6 +577,14 @@ mod tests {
error: None,
};
assert_eq!(
resolve_core_sync_error_finalize_report_kind(
"openai_memories_sync",
&result,
Some(&serde_json::json!({"error":{"message":"synthetic"}}))
),
Some("openai_memories_sync_finalize".to_string())
);
for body_json in [
serde_json::json!({"status": "failed", "error": null}),
serde_json::json!({"type": "error"}),
@@ -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);
@@ -4576,6 +4577,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(
@@ -12412,6 +12463,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 {
@@ -12436,6 +12491,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]
@@ -12631,6 +12690,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]
@@ -12870,6 +12972,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()),
@@ -45,6 +45,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
| "/v1/rerank"
| "/v1/responses"
| "/v1/responses/compact"
| "/v1/memories/trace_summarize"
| "/v1/realtime"
| "/v1/realtime/calls"
| "/v1/live"
@@ -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,8 @@
use super::super::resolve_usage_user_group_scope;
use super::range::{build_comparison_range, parse_bounded_u32};
use super::range::{
build_comparison_range, parse_bounded_u32, precise_admin_stats_time_range,
resolve_precise_time_bounds,
};
use super::resolve_admin_usage_time_range;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{
@@ -285,12 +288,36 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let time_range = match resolve_admin_usage_time_range(request_context.query_string()) {
let legacy_time_range = match resolve_admin_usage_time_range(request_context.query_string())
{
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
if let Err(detail) = time_range.validate_for_time_series(granularity) {
return Ok(Some(admin_stats_bad_request_response(detail)));
let precise_bounds = match resolve_precise_time_bounds(request_context.query_string()) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let precise_time_range = match precise_bounds {
Some((from, to)) => {
match precise_admin_stats_time_range(request_context.query_string(), from, to) {
Ok(value) => Some(value),
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
}
}
None => None,
};
let time_range = precise_time_range.as_ref().unwrap_or(&legacy_time_range);
if precise_bounds.is_none() {
if let Err(detail) = time_range.validate_for_time_series(granularity) {
return Ok(Some(admin_stats_bad_request_response(detail)));
}
} else if precise_bounds
.and_then(|(from, to)| to.checked_sub(from))
.is_some_and(|seconds| seconds > 90 * 86_400)
{
return Ok(Some(admin_stats_bad_request_response(
"Query range cannot exceed 90 days".to_string(),
)));
}
if !state.has_usage_data_reader() {
return Ok(Some(admin_stats_time_series_empty_response()));
@@ -314,7 +341,8 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
| AdminStatsGranularity::Week
| AdminStatsGranularity::Month => UsageTimeSeriesGranularity::Day,
};
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
let Some((created_from_unix_secs, created_until_unix_secs)) =
precise_bounds.or_else(|| time_range.to_unix_bounds())
else {
return Ok(Some(admin_stats_time_series_empty_response()));
};
@@ -336,7 +364,7 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
})
.await?;
return Ok(Some(build_admin_stats_time_series_response_from_summaries(
&time_range,
time_range,
granularity,
&buckets,
)));
@@ -5,7 +5,7 @@ use super::leaderboard::{
build_user_leaderboard_items_from_summaries, compare_leaderboard_items,
load_user_leaderboard_metadata, AdminStatsLeaderboardItem, AdminStatsLeaderboardNameMode,
};
use super::range::{parse_bounded_u32, parse_nonnegative_usize};
use super::range::{parse_bounded_u32, parse_nonnegative_usize, resolve_precise_time_bounds};
use super::resolve_admin_usage_time_range;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{query_param_bool, query_param_value};
@@ -228,6 +228,10 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let precise_bounds = match resolve_precise_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let metric = match AdminStatsLeaderboardMetric::parse(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
@@ -272,7 +276,8 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
"user_id is not supported for the user group leaderboard".to_string(),
)));
}
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
let Some((created_from_unix_secs, created_until_unix_secs)) =
precise_bounds.or_else(|| time_range.to_unix_bounds())
else {
return Ok(Some(build_admin_stats_user_group_leaderboard_response(
metric,
@@ -402,6 +407,10 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let precise_bounds = match resolve_precise_time_bounds(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let metric = match AdminStatsLeaderboardMetric::parse(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
@@ -442,7 +451,8 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
Ok(value) => value,
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
let Some((created_from_unix_secs, created_until_unix_secs)) =
precise_bounds.or_else(|| time_range.to_unix_bounds())
else {
return Ok(Some(admin_stats_leaderboard_empty_response(
metric,
@@ -8,7 +8,10 @@ mod leaderboard;
mod leaderboard_routes;
mod provider_quota_routes;
mod range;
pub(crate) use self::range::{parse_bounded_u32, resolve_admin_usage_time_range};
pub(crate) use self::range::{
parse_bounded_u32, precise_admin_stats_time_range, resolve_admin_usage_time_range,
resolve_precise_time_bounds, resolve_usage_time_bounds,
};
pub(crate) use aether_admin::observability::stats::{
admin_stats_bad_request_response, aggregate_usage_stats, round_to, AdminStatsTimeRange,
AdminStatsUsageFilter,
@@ -4,10 +4,14 @@ pub(super) use aether_admin::observability::stats::{
admin_usage_default_days, build_comparison_range, build_time_range_from_days, parse_naive_date,
parse_nonnegative_usize, parse_tz_offset_minutes, resolve_preset_dates, user_today,
};
use chrono::{DateTime, Offset, TimeZone, Utc};
pub(crate) fn resolve_admin_usage_time_range(
query: Option<&str>,
) -> Result<AdminStatsTimeRange, String> {
if let Some((from, to)) = resolve_precise_time_bounds(query)? {
return precise_admin_stats_time_range(query, from, to);
}
match AdminStatsTimeRange::resolve_optional(query)? {
Some(time_range) => Ok(time_range),
None => {
@@ -20,3 +24,153 @@ pub(crate) fn resolve_admin_usage_time_range(
}
}
}
/// Resolve an exact UTC range supplied by the shared admin range picker.
///
/// The older stats handlers use `start_date`/`end_date` and fixed offsets. Keep
/// that parser intact and only opt into this path when both RFC 3339 endpoints
/// are present, so existing callers retain their behavior.
pub(crate) fn resolve_precise_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 Ok(None);
}
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 = 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))
}
/// Return the exact range when present, otherwise preserve the legacy stats
/// date/preset behavior.
pub(crate) fn resolve_usage_time_bounds(query: Option<&str>) -> Result<Option<(u64, u64)>, String> {
if let Some(bounds) = resolve_precise_time_bounds(query)? {
return Ok(Some(bounds));
}
Ok(resolve_admin_usage_time_range(query)?.to_unix_bounds())
}
/// Build the date metadata used by the existing stats response builders for an
/// exact range. The data query still uses the exact UTC bounds; this metadata
/// only supplies the local date labels and offset expected by old clients.
pub(crate) fn precise_admin_stats_time_range(
query: Option<&str>,
from: u64,
to: u64,
) -> Result<AdminStatsTimeRange, String> {
let timezone_name = query_param_value(query, "timezone");
let (start_date, end_date, tz_offset_minutes) = if let Some(name) = timezone_name {
let timezone = name
.parse::<chrono_tz::Tz>()
.map_err(|_| "invalid timezone".to_string())?;
let start = Utc
.timestamp_opt(
i64::try_from(from).map_err(|_| "invalid from timestamp")?,
0,
)
.single()
.ok_or_else(|| "invalid from timestamp".to_string())?
.with_timezone(&timezone);
let end = Utc
.timestamp_opt(
i64::try_from(to.saturating_sub(1)).map_err(|_| "invalid to timestamp")?,
0,
)
.single()
.ok_or_else(|| "invalid to timestamp".to_string())?
.with_timezone(&timezone);
(
start.date_naive(),
end.date_naive(),
start.offset().fix().local_minus_utc() / 60,
)
} else {
let offset = parse_tz_offset_minutes(query)?;
let fixed = chrono::FixedOffset::east_opt(offset * 60)
.ok_or_else(|| "invalid timezone offset".to_string())?;
let start = Utc
.timestamp_opt(
i64::try_from(from).map_err(|_| "invalid from timestamp")?,
0,
)
.single()
.ok_or_else(|| "invalid from timestamp".to_string())?
.with_timezone(&fixed);
let end = Utc
.timestamp_opt(
i64::try_from(to.saturating_sub(1)).map_err(|_| "invalid to timestamp")?,
0,
)
.single()
.ok_or_else(|| "invalid to timestamp".to_string())?
.with_timezone(&fixed);
(start.date_naive(), end.date_naive(), offset)
};
Ok(AdminStatsTimeRange {
start_date,
end_date,
tz_offset_minutes,
})
}
fn query_param_value(query: Option<&str>, key: &str) -> Option<String> {
url::form_urlencoded::parse(query.unwrap_or_default().as_bytes())
.find(|(name, _)| name == key)
.map(|(_, value)| value.into_owned())
}
#[cfg(test)]
mod tests {
use super::{precise_admin_stats_time_range, resolve_precise_time_bounds};
#[test]
fn precise_stats_range_preserves_subday_bounds_and_timezone_labels() {
let query = "from=2026-09-01T23:45:00Z&to=2026-09-02T00:15:00Z&timezone=Asia%2FShanghai";
let (from, to) = resolve_precise_time_bounds(Some(query)).unwrap().unwrap();
assert_eq!(to - from, 30 * 60);
let range = precise_admin_stats_time_range(Some(query), from, to).unwrap();
assert_eq!(range.start_date.to_string(), "2026-09-02");
assert_eq!(range.end_date.to_string(), "2026-09-02");
assert_eq!(range.tz_offset_minutes, 480);
}
#[test]
fn precise_stats_range_rejects_mixed_legacy_presets() {
let query = "from=2026-09-01T00:00:00Z&to=2026-09-02T00:00:00Z&preset=today";
assert!(resolve_precise_time_bounds(Some(query)).is_err());
}
}
@@ -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() {
@@ -1,5 +1,5 @@
use super::super::resolve_usage_user_group_scope;
use super::super::stats::resolve_admin_usage_time_range;
use super::super::stats::resolve_usage_time_bounds;
use super::analytics::admin_usage_api_key_names;
use super::analytics::admin_usage_provider_key_names;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
@@ -34,6 +34,12 @@ 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> {
resolve_usage_time_bounds(query)
}
async fn load_admin_usage_by_ids(
state: &AdminAppState<'_>,
requested_ids: &BTreeSet<String>,
@@ -71,6 +77,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()]);
}
@@ -565,12 +572,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,
@@ -578,11 +610,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> {
@@ -695,6 +739,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(),
@@ -741,13 +794,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(),
)));
@@ -792,13 +843,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(),
@@ -855,7 +904,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))),
};
@@ -876,9 +925,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(),
@@ -898,13 +945,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() {
@@ -1077,6 +1127,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,
};
@@ -1358,7 +1491,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);
@@ -1379,7 +1512,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
{
@@ -2,6 +2,7 @@
use async_trait::async_trait;
use serde_json::Value;
use std::borrow::Cow;
use super::adapters::CODEX_RESPONSES_WEBSOCKET_ADAPTER;
use crate::ai_serving::AiExecutionDecision;
@@ -59,7 +60,7 @@ pub(super) enum ResponsesWebSocketRelayDirective<'a> {
ForwardOriginal,
/// The provider frame was a private batch envelope. Forward each retained
/// event in document order by serializing the complete borrowed value.
ForwardEvents(Vec<&'a Value>),
ForwardEvents(Vec<Cow<'a, Value>>),
/// The entire frame was an explicitly recognized provider-private
/// envelope and therefore has no public event to relay.
SuppressProviderPrivate,
@@ -88,8 +89,8 @@ pub(super) trait ResponsesWebSocketProtocolAdapter: Send + Sync {
/// observably ambiguous to the client.
fn rebind_safety_for_upstream_event(&self, event: &Value) -> ResponsesWebSocketRebindSafety;
/// Selects the public relay shape without projecting a provider event
/// through an Aether-owned field or event-type allowlist.
/// 公开事件保持完整;提供商私有元数据仅投影到客户端使用的公开字段,
/// 不公开账户信息。
fn relay_directive_for_upstream_event<'a>(
&self,
_event: &'a Value,
@@ -2,6 +2,7 @@
use async_trait::async_trait;
use serde_json::{Map, Value};
use std::borrow::Cow;
use super::super::adapter::{
is_standard_responses_event, ResponsesWebSocketAdapterObservation,
@@ -222,7 +223,7 @@ fn codex_relay_directive(event: &Value) -> ResponsesWebSocketRelayDirective<'_>
Some(Value::Array(chunks)) if is_explicit_codex_batch_envelope(event) => {
let public_events = chunks
.iter()
.filter(|chunk| !is_codex_private_leaf_event(chunk))
.filter_map(codex_public_event)
.collect::<Vec<_>>();
if public_events.is_empty() {
ResponsesWebSocketRelayDirective::SuppressProviderPrivate
@@ -233,13 +234,51 @@ fn codex_relay_directive(event: &Value) -> ResponsesWebSocketRelayDirective<'_>
// A malformed or future shape is not proven private. Preserve it
// opaquely rather than guessing at a provider schema.
Some(_) => ResponsesWebSocketRelayDirective::ForwardOriginal,
None if is_codex_private_leaf_event(event) => {
ResponsesWebSocketRelayDirective::SuppressProviderPrivate
}
None if is_codex_private_leaf_event(event) => match codex_public_event(event) {
Some(projected) => ResponsesWebSocketRelayDirective::ForwardEvents(vec![projected]),
None => ResponsesWebSocketRelayDirective::SuppressProviderPrivate,
},
None => ResponsesWebSocketRelayDirective::ForwardOriginal,
}
}
fn codex_public_event(event: &Value) -> Option<Cow<'_, Value>> {
if !is_codex_private_leaf_event(event) {
return Some(Cow::Borrowed(event));
}
if event.get("type").and_then(Value::as_str) != Some("codex.response.metadata") {
// 配额属于选中的上游账户,不能代表网关用户配额;仅交给
// 账户级熔断和持久化路径处理。
return None;
}
let headers = event.get("headers")?.as_object()?;
let public_headers: Map<String, Value> = headers
.iter()
.filter_map(|(name, value)| {
let name = name.to_ascii_lowercase();
if matches!(
name.as_str(),
"x-models-etag"
| "x-codex-turn-state"
| "openai-model"
| "x-codex-safety-buffering-enabled"
| "x-codex-safety-buffering-faster-model"
) && value.as_str().is_some()
{
Some((name, value.clone()))
} else {
None
}
})
.collect();
if public_headers.is_empty() {
return None;
}
Some(Cow::Owned(serde_json::json!({
"type": "codex.response.metadata", "headers": public_headers
})))
}
/// Recognizes only Codex's private batch container. A type-less object must
/// contain exactly `chunks`; unknown siblings could be future public protocol
/// data and therefore force opaque forwarding. A named Codex private root may
@@ -549,6 +588,55 @@ mod tests {
}
}
#[test]
fn codex_metadata_relays_cli_catalog_and_turn_state_without_account_fields() {
let event = json!({
"type": "codex.response.metadata",
"headers": {
"X-Models-Etag": "catalog-v2",
"x-codex-turn-state": "synthetic-turn-state",
"openai-model": "gpt-6.1-sol",
"x-codex-safety-buffering-enabled": "true",
"x-codex-safety-buffering-faster-model": "gpt-6-luna",
"chatgpt-account-id": "synthetic-private-account",
"set-cookie": "synthetic-private-cookie",
"authorization": "synthetic-private-token"
},
"account_hint": "private",
"metadata": {"user_id": "private"}
});
let ResponsesWebSocketRelayDirective::ForwardEvents(events) =
CodexResponsesWebSocketAdapter.relay_directive_for_upstream_event(&event)
else {
panic!("CLI metadata must reach the client");
};
assert_eq!(events.len(), 1);
assert_eq!(
*events[0],
json!({
"type": "codex.response.metadata",
"headers": {"x-models-etag": "catalog-v2", "x-codex-turn-state": "synthetic-turn-state", "openai-model": "gpt-6.1-sol", "x-codex-safety-buffering-enabled": "true", "x-codex-safety-buffering-faster-model": "gpt-6-luna"}
})
);
}
#[test]
fn codex_batch_retains_safe_metadata_in_public_event_order() {
let event = json!({"chunks": [
{"type": "codex.response.metadata", "headers": {"x-models-etag": "catalog-v3"}},
{"type": "codex.rate_limits", "plan_type": "private-plan"},
{"type": "response.created", "response": {"id": "resp_synthetic"}, "future": 42}
]});
let ResponsesWebSocketRelayDirective::ForwardEvents(events) =
CodexResponsesWebSocketAdapter.relay_directive_for_upstream_event(&event)
else {
panic!("batch must retain public events");
};
assert_eq!(events.len(), 2);
assert_eq!(events[0]["headers"]["x-models-etag"], "catalog-v3");
assert_eq!(events[1]["future"], 42);
}
#[test]
fn mixed_codex_batch_forwards_whole_non_private_events_in_order() {
let adapter = CodexResponsesWebSocketAdapter;
@@ -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);
@@ -610,6 +610,7 @@ pub(super) async fn relay_bound_connection(
}
Some(ResponsesWebSocketRelayDirective::ForwardEvents(events)) => {
for event in events {
let event = event.as_ref();
let text = match bound
.redaction_restorer
.restore_provider_frame_text(event)
@@ -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!(
@@ -1211,6 +1211,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)
@@ -1261,6 +1262,7 @@ pub(super) async fn handle_users_me_usage_get(
limit: None,
offset: None,
newest_first: true,
..Default::default()
})
.await
{
@@ -1291,6 +1293,7 @@ pub(super) async fn handle_users_me_usage_get(
limit: Some(limit),
offset: Some(offset),
newest_first: true,
..Default::default()
})
.await
{
@@ -1436,6 +1439,7 @@ pub(super) async fn handle_users_me_usage_active_get(
limit: Some(50),
offset: None,
newest_first: true,
..Default::default()
})
.await
{
@@ -1926,6 +1930,29 @@ mod tests {
assert_eq!(active["reasoning_effort"], "max");
}
#[test]
fn user_usage_payloads_expose_gemini_thinking_config_reasoning_mapping() {
let item = StoredRequestUsageAudit {
request_body: Some(json!({
"generationConfig": {
"thinkingConfig": { "includeThoughts": true, "thinkingLevel": "HIGH" }
}
})),
provider_request_body: Some(json!({
"generationConfig": { "thinkingConfig": { "thinkingBudget": 8192 } }
})),
..sample_usage("completed")
};
let record = build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
let active = build_users_me_usage_active_payload(&item);
assert_eq!(record["requested_reasoning_effort"], "high");
assert_eq!(active["requested_reasoning_effort"], "high");
assert_eq!(record["reasoning_effort"], "xhigh");
assert_eq!(active["reasoning_effort"], "xhigh");
}
#[test]
fn user_usage_payloads_expose_websocket_transport() {
let item = StoredRequestUsageAudit {
@@ -2681,6 +2681,84 @@ fn windsurf_quota_snapshot_has_stale_cooldown(quota_snapshot: &Map<String, Value
has_capacity || !exhausted
}
/// 读取状态快照时,把已越过重置时间点的配额窗口按“已重置”口径归一化。
///
/// 背景:调度侧早已把到期窗口视为未耗尽(`provider_pool_reset_deadline_elapsed`),
/// 账号额度文本也会按到期强制显示 100%,但列表读取此前直接返回存量快照,
/// 导致管理端倒计时归零后进度条仍停留在旧的剩余百分比。这里让读取层与
/// 调度侧、文本侧使用同一口径,避免三处状态互相矛盾。
///
/// 覆盖所有带重置时间的提供商窗口(codex/kiro/xai/grok/antigravity/
/// gemini_cli/chatgpt_web/windsurf 等):比例(used_ratio/remaining_ratio)、
/// 数值(used_value/remaining_value)与窗口级耗尽标记会一起恢复为“已重置”。
fn normalize_expired_quota_windows(snapshot: &mut serde_json::Map<String, Value>) {
let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) else {
return;
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let fallback_observed_at = provider_quota_timestamp_unix_secs(quota.get("observed_at"))
.or_else(|| provider_quota_timestamp_unix_secs(quota.get("updated_at")));
let Some(windows) = quota.get_mut("windows").and_then(Value::as_array_mut) else {
return;
};
for window in windows.iter_mut().filter_map(Value::as_object_mut) {
// 与同类窗口处理保持一致:window_minutes=0 不是真实配额窗口。
if window.get("window_minutes").and_then(Value::as_u64) == Some(0) {
continue;
}
if !aether_provider_pool::provider_pool_reset_deadline_elapsed(
window,
fallback_observed_at,
now_unix_secs,
) {
continue;
}
// 只归一化“观测完整”的窗口:要么有比例观测,要么有上限 + 用量数值。
// 这样既不会把无数据窗口凭空显示成 100%,也不会出现比例已恢复 100%
// 而数值仍停留在旧值的不一致(例如只有 remaining_value 却没有上限的窗口)。
let has_ratio_observation = ["used_ratio", "remaining_ratio"]
.into_iter()
.any(|field| window.get(field).is_some_and(Value::is_number));
let limit_value = window
.get("limit_value")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| *value > 0.0);
let has_value_observation = limit_value.is_some()
&& ["used_value", "remaining_value"]
.into_iter()
.any(|field| window.get(field).is_some_and(Value::is_number));
if !has_ratio_observation && !has_value_observation {
continue;
}
// 比例口径:已用清零、剩余 100%;字段原本为 null 时一并补齐,保证展示口径统一。
for (field, value) in [("used_ratio", 0.0), ("remaining_ratio", 1.0)] {
if window.contains_key(field) {
window.insert(field.to_string(), json!(value));
}
}
// 数值口径:已用清零;有上限时把剩余恢复到上限(用于“剩余 x/y”类文本展示)。
if window.contains_key("used_value") {
window.insert("used_value".to_string(), json!(0.0));
}
if window.contains_key("remaining_value") {
if let Some(limit_value) = limit_value {
window.insert("remaining_value".to_string(), json!(limit_value));
}
}
// 同步清掉窗口级耗尽标记,避免展示与调度口径互相矛盾。
for field in ["is_exhausted", "exhausted"] {
if let Some(slot) = window.get_mut(field) {
*slot = json!(false);
}
}
}
}
pub(crate) fn provider_key_status_snapshot_payload(
key: &StoredProviderCatalogKey,
provider_type: &str,
@@ -2722,9 +2800,8 @@ pub(crate) fn provider_key_status_snapshot_payload(
// Legacy snapshots can retain an exhausted summary after a window reset or
// newer quota observation. Use the same decision as scheduling so the
// account list and its status filter do not keep displaying that stale block.
if provider_type.trim().eq_ignore_ascii_case("codex")
&& !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type)
{
// 对所有提供商生效:适配器判定已是“重置感知”的,与调度口径保持一致。
if !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type) {
if let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) {
quota.insert("exhausted".to_string(), json!(false));
if quota.get("code").and_then(Value::as_str) == Some("exhausted") {
@@ -2732,6 +2809,9 @@ pub(crate) fn provider_key_status_snapshot_payload(
}
}
}
// 读取时归一化已到期的配额窗口,保证列表进度条、额度文字与调度侧、
// 账号额度文本使用同一“已重置”口径。
normalize_expired_quota_windows(&mut snapshot);
snapshot.insert(
"oauth".to_string(),
build_provider_key_oauth_status_snapshot(key),
@@ -3780,6 +3860,421 @@ mod tests {
assert_eq!(window.get("reset_seconds"), Some(&json!(3_600u64)));
}
#[test]
fn provider_key_status_snapshot_payload_normalizes_expired_codex_quota_windows() {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"exhausted": false,
"updated_at": 1_700_000_000u64,
"windows": [
{
"code": "5h",
"label": "5H",
"scope": "account",
"unit": "percent",
"used_ratio": 0.88,
"remaining_ratio": 0.12,
"reset_at": 1_700_003_600u64,
"window_minutes": 300
},
{
"code": "weekly",
"label": "周",
"scope": "account",
"unit": "percent",
"used_ratio": 0.5,
"remaining_ratio": 0.5,
"reset_at": 2_000_000_000u64,
"window_minutes": 10_080
},
{
"code": "spark_5h",
"label": "Spark 5H",
"scope": "account",
"unit": "percent",
"reset_at": 1_700_003_600u64,
"window_minutes": 300
},
{
"code": "unlimited",
"label": "无限",
"scope": "account",
"unit": "percent",
"used_ratio": 0.3,
"remaining_ratio": 0.7,
"reset_at": 1_700_003_600u64,
"window_minutes": 0
}
]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
// 已到期的窗口按“已重置”口径归一化:用量清零、剩余 100%。
assert_eq!(payload.pointer("/quota/windows/0/code"), Some(&json!("5h")));
assert_eq!(
payload.pointer("/quota/windows/0/used_ratio"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_ratio"),
Some(&json!(1.0))
);
// 未到期的窗口保留原始观测值。
assert_eq!(
payload.pointer("/quota/windows/1/code"),
Some(&json!("weekly"))
);
assert_eq!(
payload.pointer("/quota/windows/1/used_ratio"),
Some(&json!(0.5))
);
assert_eq!(
payload.pointer("/quota/windows/1/remaining_ratio"),
Some(&json!(0.5))
);
// 没有用量观测的窗口不会被凭空补成 100%。
assert_eq!(
payload.pointer("/quota/windows/2/code"),
Some(&json!("spark_5h"))
);
assert_eq!(payload.pointer("/quota/windows/2/used_ratio"), None);
assert_eq!(payload.pointer("/quota/windows/2/remaining_ratio"), None);
// window_minutes=0 不是真实配额窗口,与同类 Codex 窗口处理保持一致,不归一化。
assert_eq!(
payload.pointer("/quota/windows/3/code"),
Some(&json!("unlimited"))
);
assert_eq!(
payload.pointer("/quota/windows/3/used_ratio"),
Some(&json!(0.3))
);
assert_eq!(
payload.pointer("/quota/windows/3/remaining_ratio"),
Some(&json!(0.7))
);
}
#[test]
fn provider_key_status_snapshot_payload_clears_expired_codex_window_exhausted_state() {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"usage_ratio": 1.0,
"updated_at": 1_700_000_000u64,
"windows": [
{
"code": "weekly",
"label": "周",
"scope": "account",
"unit": "percent",
"used_ratio": 1.0,
"remaining_ratio": 0.0,
"reset_at": 1_700_003_600u64,
"window_minutes": 10_080,
"is_exhausted": true
}
]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
// 汇总标志沿用既有调度口径(窗口已到期 → 不再视为耗尽)。
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(false)));
assert_eq!(payload.pointer("/quota/code"), Some(&json!("ok")));
// 窗口级耗尽标记与剩余比例同步归一化。
assert_eq!(
payload.pointer("/quota/windows/0/is_exhausted"),
Some(&json!(false))
);
assert_eq!(
payload.pointer("/quota/windows/0/used_ratio"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_ratio"),
Some(&json!(1.0))
);
}
#[test]
fn provider_key_status_snapshot_payload_normalizes_expired_quota_windows_for_other_providers() {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "kiro",
"code": "ok",
"exhausted": false,
"observed_at": 1_700_000_000u64,
"updated_at": 1_700_000_000u64,
"windows": [
{
"code": "usage",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 0.4,
"remaining_ratio": 0.6,
"used_value": 60.0,
"remaining_value": 90.0,
"limit_value": 150.0,
"reset_at": 1_700_003_600u64,
"reset_seconds": 3_600u64
},
{
"code": "usage_null_ratio",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": null,
"remaining_ratio": null,
"used_value": 60.0,
"remaining_value": 90.0,
"limit_value": 150.0,
"reset_at": 1_700_003_600u64,
"reset_seconds": 3_600u64
},
{
"code": "usage_active",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 0.4,
"remaining_ratio": 0.6,
"used_value": 60.0,
"remaining_value": 90.0,
"limit_value": 150.0,
"reset_at": 2_000_000_000u64,
"reset_seconds": 300_000_000u64
},
{
"code": "usage_no_observation",
"label": "额度",
"scope": "account",
"unit": "count",
"reset_at": 1_700_003_600u64,
"reset_seconds": 3_600u64
},
{
"code": "usage_no_deadline",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 0.4,
"remaining_ratio": 0.6,
"used_value": 60.0,
"remaining_value": 90.0,
"limit_value": 150.0
},
{
"code": "usage_remaining_without_limit",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": null,
"remaining_ratio": null,
"remaining_value": 0.0,
"reset_at": 1_700_003_600u64,
"reset_seconds": 3_600u64
}
]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "kiro");
// 已到期窗口:比例与数值一起按“已重置”归一化(不再只覆盖 Codex)。
assert_eq!(
payload.pointer("/quota/windows/0/used_ratio"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_ratio"),
Some(&json!(1.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/used_value"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_value"),
Some(&json!(150.0))
);
// 比例字段为 null 时补齐为已重置口径,保证优先读比例的展示实现一致。
assert_eq!(
payload.pointer("/quota/windows/1/used_ratio"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/1/remaining_ratio"),
Some(&json!(1.0))
);
assert_eq!(
payload.pointer("/quota/windows/1/used_value"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/1/remaining_value"),
Some(&json!(150.0))
);
// 未到期窗口保留原始观测值。
assert_eq!(
payload.pointer("/quota/windows/2/used_ratio"),
Some(&json!(0.4))
);
assert_eq!(
payload.pointer("/quota/windows/2/remaining_value"),
Some(&json!(90.0))
);
// 没有用量观测的窗口不会被凭空补成 100%。
assert_eq!(payload.pointer("/quota/windows/3/used_ratio"), None);
assert_eq!(payload.pointer("/quota/windows/3/remaining_ratio"), None);
// 没有重置时间的窗口无法判定是否已重置,保持原样。
assert_eq!(
payload.pointer("/quota/windows/4/used_ratio"),
Some(&json!(0.4))
);
assert_eq!(
payload.pointer("/quota/windows/4/remaining_ratio"),
Some(&json!(0.6))
);
// 只有 remaining_value 却没有上限的窗口观测不完整,跳过归一化,
// 避免比例已恢复 100% 而数值仍停留在旧值。
assert_eq!(
payload.pointer("/quota/windows/5/remaining_ratio"),
Some(&json!(null))
);
assert_eq!(
payload.pointer("/quota/windows/5/remaining_value"),
Some(&json!(0.0))
);
}
#[test]
fn provider_key_status_snapshot_payload_clears_expired_quota_exhausted_state_for_other_providers(
) {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "kiro",
"code": "exhausted",
"label": "额度耗尽",
"reason": "额度已耗尽",
"exhausted": true,
"usage_ratio": 1.0,
"observed_at": 1_700_000_000u64,
"updated_at": 1_700_000_000u64,
"windows": [
{
"code": "usage",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 1.0,
"remaining_ratio": 0.0,
"used_value": 150.0,
"remaining_value": 0.0,
"limit_value": 150.0,
"reset_at": 1_700_003_600u64,
"reset_seconds": 3_600u64,
"is_exhausted": true
}
]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "kiro");
// 汇总状态沿用调度口径(窗口已到期 → 不再视为耗尽),对所有提供商生效。
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(false)));
assert_eq!(payload.pointer("/quota/code"), Some(&json!("ok")));
// 窗口比例、数值与耗尽标记同步归一化。
assert_eq!(
payload.pointer("/quota/windows/0/is_exhausted"),
Some(&json!(false))
);
assert_eq!(
payload.pointer("/quota/windows/0/used_ratio"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_ratio"),
Some(&json!(1.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/used_value"),
Some(&json!(0.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_value"),
Some(&json!(150.0))
);
}
#[test]
fn provider_key_status_snapshot_payload_keeps_active_quota_exhausted_state_for_other_providers()
{
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "kiro",
"code": "exhausted",
"label": "额度耗尽",
"exhausted": true,
"usage_ratio": 1.0,
"observed_at": 1_700_000_000u64,
"updated_at": 1_700_000_000u64,
"windows": [
{
"code": "usage",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 1.0,
"remaining_ratio": 0.0,
"used_value": 150.0,
"remaining_value": 0.0,
"limit_value": 150.0,
"reset_at": 2_000_000_000u64,
"reset_seconds": 300_000_000u64,
"is_exhausted": true
}
]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "kiro");
// 未到期的耗尽状态保留,不能被读取层提前“恢复”。
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(true)));
assert_eq!(payload.pointer("/quota/code"), Some(&json!("exhausted")));
assert_eq!(
payload.pointer("/quota/windows/0/is_exhausted"),
Some(&json!(true))
);
assert_eq!(
payload.pointer("/quota/windows/0/used_ratio"),
Some(&json!(1.0))
);
assert_eq!(
payload.pointer("/quota/windows/0/remaining_ratio"),
Some(&json!(0.0))
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_chatgpt_web_image_quota() {
let mut key = sample_catalog_key();
@@ -3911,14 +4406,14 @@ mod tests {
"used_percent": 60.0,
"remaining": 60.0,
"total": 150.0,
"reset_at": 1_778_157_172u64,
"reset_at": 2_000_000_000u64,
"is_exhausted": false
},
"quota_heavy": {
"display_name": "heavy",
"remaining_fraction": 0.0,
"used_percent": 100.0,
"reset_at": 1_778_157_172u64,
"reset_at": 2_000_000_000u64,
"is_exhausted": true
}
}
@@ -3941,7 +4436,7 @@ mod tests {
assert_eq!(quota.get("pool_tier"), Some(&json!("heavy")));
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
assert_eq!(quota.get("reset_at"), Some(&json!(2_000_000_000u64)));
assert_eq!(windows.len(), 2);
assert!(windows.iter().any(|window| {
window
@@ -4054,8 +4549,8 @@ mod tests {
"plan_name": "Pro",
"daily_remaining_percent": 40.0,
"weekly_remaining_percent": 65.0,
"daily_reset_at": 1_778_100_000u64,
"weekly_reset_at": 1_778_600_000u64,
"daily_reset_at": 2_000_000_000u64,
"weekly_reset_at": 2_000_600_000u64,
"prompt_used": 12.0,
"prompt_limit": 100.0,
"prompt_remaining": 88.0,
@@ -4090,13 +4585,13 @@ mod tests {
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("plan_type"), Some(&json!("Pro")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.6)));
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_100_000u64)));
assert_eq!(quota.get("reset_at"), Some(&json!(2_000_000_000u64)));
assert_eq!(daily.get("remaining_ratio"), Some(&json!(0.4)));
assert_eq!(daily.get("used_ratio"), Some(&json!(0.6)));
assert_eq!(daily.get("reset_seconds"), Some(&json!(32_754u64)));
assert_eq!(daily.get("reset_seconds"), Some(&json!(221_932_754u64)));
assert_eq!(weekly.get("remaining_ratio"), Some(&json!(0.65)));
assert_eq!(weekly.get("used_ratio"), Some(&json!(0.35)));
assert_eq!(weekly.get("reset_seconds"), Some(&json!(532_754u64)));
assert_eq!(weekly.get("reset_seconds"), Some(&json!(222_532_754u64)));
assert_eq!(quota.get("allowed_models_count"), Some(&json!(82)));
}
@@ -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 / 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"))
@@ -345,6 +346,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"))
+3
View File
@@ -44,6 +44,7 @@ mod data;
mod dispatch;
mod email_delivery;
mod error;
mod execution_activity;
mod execution_runtime;
mod executor;
mod fallback_metrics;
@@ -69,6 +70,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;
@@ -90,6 +92,7 @@ mod upstream_admission;
mod usage;
mod video_tasks;
mod wallet_runtime;
mod xai_profile;
pub use self::ai_serving::api::{codex_client_originator, codex_client_user_agent};
pub(crate) use self::ai_serving::api::{
+14
View File
@@ -2546,6 +2546,20 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
);
}
}
match state.prewarm_xai_client_profile().await {
Ok(version) => {
info!(
xai_client_version = %version,
"prewarmed Grok CLI client profile"
);
}
Err(err) => {
warn!(
error = %err,
"failed to refresh Grok CLI client profile; built-in or cached profile remains active"
);
}
}
match prewarm_direct_h2c_sender_cache_from_env_for_startup().await {
Ok(Some(report)) => {
if report.failed_targets > 0 {
@@ -155,14 +155,17 @@ pub(crate) async fn perform_oauth_token_refresh_once(
Ok(None) => {
summary.skipped = summary.skipped.saturating_add(1);
}
Err(_) => {
Err(err) => {
summary.failed = summary.failed.saturating_add(1);
warn!(
event_name = "oauth_token_refresh_failed",
log_type = "ops",
worker = "oauth_token_refresh",
provider_id = %provider.id,
provider_type = %provider.provider_type,
endpoint_id = %endpoint.id,
key_id = %key.id,
error = %crate::error::redact_error_debug(&err),
"gateway oauth token auto refresh failed"
);
}
@@ -440,6 +443,7 @@ mod tests {
agent_identity_needs_task_recovery, auth_config_has_refresh_token,
is_nonfatal_legacy_catalog_credential_error, oauth_refresh_candidate,
};
use crate::error::redact_error_debug;
use crate::GatewayError;
#[test]
@@ -543,4 +547,18 @@ mod tests {
)
));
}
#[test]
fn oauth_refresh_failure_detail_preserves_context_without_credentials() {
let error = GatewayError::Internal(
r#"oauth request failed: status=503 token="refresh-secret" retry=2"#.to_string(),
);
let detail = redact_error_debug(&error);
assert!(detail.contains("oauth request failed"));
assert!(detail.contains("status=503"));
assert!(detail.contains("[REDACTED]"));
assert!(!detail.contains("refresh-secret"));
}
}
@@ -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);
},
))
}
@@ -2663,14 +2663,11 @@ mod tests {
assert_eq!(plans.len(), 2);
assert!(plans[0].url.ends_with("/models?client_version=0.145.2"));
assert!(plans[1].url.ends_with("/models?client_version=0.146.0"));
assert_eq!(
plans[0].headers.get("user-agent").map(String::as_str),
Some("codex_cli_rs/0.145.2")
);
assert_eq!(
plans[1].headers.get("user-agent").map(String::as_str),
Some("codex_cli_rs/0.146.0")
);
for (plan, version) in plans.iter().zip(["0.145.2", "0.146.0"]) {
let ua = plan.headers.get("user-agent").expect("画像应包含 UA");
assert!(ua.starts_with(&format!("codex_cli_rs/{version} (")));
assert!(ua.ends_with(") unknown"));
}
let target = target();
assert_ne!(
@@ -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>,
+16 -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;
@@ -77,6 +78,7 @@ use crate::maintenance::spawn_stats_hourly_aggregation_worker;
use crate::maintenance::spawn_usage_cleanup_worker;
use crate::maintenance::spawn_usage_counter_flush_worker;
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
use crate::xai_profile::spawn_worker as spawn_xai_client_profile_worker;
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(30);
// Requests may use a stale value after the fresh window until the entry reaches
@@ -159,6 +161,10 @@ impl AppState {
crate::cli_client_profile::prewarm(&CLAUDE_CODE_CLI_PROFILE, self.runtime_state()).await
}
pub async fn prewarm_xai_client_profile(&self) -> Result<String, String> {
crate::xai_profile::prewarm(self.runtime_state()).await
}
pub async fn prewarm_chat_pii_redaction_runtime_config(&self) -> Result<bool, String> {
crate::privacy::read_chat_pii_redaction_runtime_config(self)
.await
@@ -265,6 +271,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();
@@ -362,6 +369,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,
)),
@@ -413,6 +422,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()),
@@ -2367,6 +2377,10 @@ impl AppState {
background_state.clone(),
)),
);
supervise_worker(
crate::task_runtime::TASK_KEY_XAI_CLIENT_PROFILE,
Some(spawn_xai_client_profile_worker(background_state.clone())),
);
supervise_worker(
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
spawn_video_task_poller(background_state.clone()),
+90 -8
View File
@@ -412,6 +412,17 @@ fn normalize_local_oauth_refresh_error_message(
.unwrap_or_else(|| "Token 刷新失败".to_string())
}
fn local_oauth_refresh_gateway_error(
error: &provider_transport::LocalOAuthRefreshError,
) -> GatewayError {
// Keep a bounded, credential-redacted reason for internal diagnostics.
// GatewayError::Internal still returns the generic error response to clients.
GatewayError::Internal(format!(
"local oauth refresh failed: {}",
crate::error::redact_error_detail(error)
))
}
fn merge_local_oauth_refresh_failure_reason(
current_reason: Option<&str>,
refresh_reason: &str,
@@ -1556,10 +1567,8 @@ impl AppState {
}
return Ok(None);
}
Err(_) => {
return Err(GatewayError::Internal(
"local oauth refresh failed".to_string(),
));
Err(err) => {
return Err(local_oauth_refresh_gateway_error(&err));
}
};
@@ -3484,13 +3493,86 @@ mod tests {
use tokio::sync::Notify;
use super::{
AgentIdentityAuthConfigFence, AppState, CodexRuntimeOAuthObservation,
ProviderTransportSnapshotCacheKey, ProviderTransportSnapshotFlight,
ProviderTransportSnapshotFlightResult, ProviderTransportSnapshotInflightRegistration,
PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL,
local_oauth_refresh_gateway_error, AgentIdentityAuthConfigFence, AppState,
CodexRuntimeOAuthObservation, ProviderTransportSnapshotCacheKey,
ProviderTransportSnapshotFlight, ProviderTransportSnapshotFlightResult,
ProviderTransportSnapshotInflightRegistration, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL,
PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL,
};
use crate::data::GatewayDataState;
#[test]
fn oauth_refresh_diagnostic_preserves_failure_reason_and_redacts_credentials() {
for (message, secret) in [
(
"connection refused refresh_token=refresh-secret",
"refresh-secret",
),
(
"connection refused accessToken=access-secret",
"access-secret",
),
(
"connection refused client_secret=client-secret",
"client-secret",
),
(
"connection refused Authorization: Bearer bearer-secret",
"bearer-secret",
),
(
"connection refused https://proxy-user:[email protected]",
"proxy-secret",
),
] {
let error = crate::provider_transport::LocalOAuthRefreshError::TransportMessage {
provider_type: "codex",
message: message.to_string(),
};
let diagnostic = local_oauth_refresh_gateway_error(&error).into_message();
assert!(diagnostic.contains("codex oauth refresh transport failed"));
assert!(diagnostic.contains("connection refused"));
assert!(!diagnostic.contains(secret), "diagnostic: {diagnostic}");
}
}
#[test]
fn oauth_refresh_diagnostic_keeps_http_status_without_provider_body() {
let error = crate::provider_transport::LocalOAuthRefreshError::HttpStatus {
provider_type: "codex",
status_code: 503,
body_excerpt: "unstructured-provider-secret".to_string(),
};
let diagnostic = local_oauth_refresh_gateway_error(&error).into_message();
assert!(diagnostic.contains("codex oauth refresh returned HTTP 503"));
assert!(!diagnostic.contains("unstructured-provider-secret"));
}
#[tokio::test]
async fn oauth_refresh_diagnostic_is_hidden_from_client_response() {
use axum::body::to_bytes;
use axum::response::IntoResponse;
let error = crate::provider_transport::LocalOAuthRefreshError::TransportMessage {
provider_type: "codex",
message: "connection refused".to_string(),
};
let response = local_oauth_refresh_gateway_error(&error).into_response();
assert_eq!(
response.status(),
axum::http::StatusCode::INTERNAL_SERVER_ERROR
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("error response body should read");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("error should be JSON"),
json!({"error": {"message": "internal server error"}}),
);
}
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
@@ -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
@@ -27,6 +27,7 @@ pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_CODEX_CLIENT_PROFILE: &str = "maintenance.codex.client.profile";
pub(crate) const TASK_KEY_CLAUDE_CODE_CLIENT_PROFILE: &str =
"maintenance.claude_code.client.profile";
pub(crate) const TASK_KEY_XAI_CLIENT_PROFILE: &str = "maintenance.xai.client.profile";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
@@ -221,6 +222,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_XAI_CLIENT_PROFILE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_QUOTA_RESET,
TaskKind::Scheduled,
@@ -958,3 +967,28 @@ mod worker_boot_run_id_tests {
assert!(events.iter().all(|event| event.payload_json.is_none()));
}
}
#[cfg(test)]
mod client_profile_task_tests {
use super::*;
#[test]
fn all_cli_client_profile_tasks_are_registered_once() {
for (key, trigger) in [
(TASK_KEY_CODEX_CLIENT_PROFILE, "daily"),
(TASK_KEY_CLAUDE_CODE_CLIENT_PROFILE, "daily"),
(TASK_KEY_XAI_CLIENT_PROFILE, "interval"),
] {
let definitions: Vec<_> = task_definitions()
.iter()
.filter(|definition| definition.key == key)
.collect();
assert_eq!(definitions.len(), 1, "task {key} must be registered once");
let definition = definitions[0];
assert_eq!(definition.kind, TaskKind::Scheduled);
assert_eq!(definition.trigger, trigger);
assert!(definition.singleton);
assert!(definition.persist_history);
}
}
}
@@ -45,6 +45,153 @@ where
}
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
"user-search-1".to_string(),
"alice".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
"api-key-search-1".to_string(),
Some("search-client".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
)
.expect("auth snapshot should build")
}
fn candidate_row(api_format: &str) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-codex-search-1".to_string(),
provider_name: "codex".to_string(),
provider_type: "codex".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-codex-search-1".to_string(),
endpoint_api_format: api_format.to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some(api_format.split_once(':').expect("format").1.to_string()),
endpoint_is_active: true,
key_id: "key-codex-search-1".to_string(),
key_name: "oauth".to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:responses".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(json!({"openai:search": 1})),
model_id: "model-codex-search-1".to_string(),
global_model_id: "global-model-codex-search-1".to_string(),
global_model_name: "gpt-5.6-sol".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gpt-5.6-sol".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5.6-sol".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex-search-1".to_string(),
"codex".to_string(),
Some("https://chatgpt.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(1),
None,
Some(900.0),
None,
None,
)
}
fn endpoint(api_format: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
api_format.to_string(),
Some("openai".to_string()),
Some(api_format.split_once(':').expect("format").1.to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(1),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn key() -> StoredProviderCatalogKey {
let auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","account_id":"account-search-1","is_fedramp":true}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
"oauth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:responses"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "codex-search-access-token")
.expect("access token should encrypt"),
Some(auth_config),
None,
Some(json!({"openai:search": 1})),
None,
Some(4_102_444_800),
None,
None,
)
.expect("key transport should build")
}
#[test]
fn gateway_executes_codex_search_with_responses_permission_and_search_contract() {
run_search_sync_test(
@@ -54,156 +201,6 @@ fn gateway_executes_codex_search_with_responses_permission_and_search_contract()
}
async fn gateway_executes_codex_search_with_responses_permission_and_search_contract_impl() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
"user-search-1".to_string(),
"alice".to_string(),
Some("[email protected]".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
"api-key-search-1".to_string(),
Some("search-client".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(json!(["openai", "codex"])),
Some(json!(["openai:responses"])),
None,
)
.expect("auth snapshot should build")
}
fn candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-codex-search-1".to_string(),
provider_name: "codex".to_string(),
provider_type: "codex".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-codex-search-1".to_string(),
endpoint_api_format: "openai:search".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("search".to_string()),
endpoint_is_active: true,
key_id: "key-codex-search-1".to_string(),
key_name: "oauth".to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:responses".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(json!({"openai:search": 1})),
model_id: "model-codex-search-1".to_string(),
global_model_id: "global-model-codex-search-1".to_string(),
global_model_name: "gpt-5.6-sol".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gpt-5.6-sol".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5.6-sol".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex-search-1".to_string(),
"codex".to_string(),
Some("https://chatgpt.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(1),
None,
Some(900.0),
None,
None,
)
}
fn endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
"openai:search".to_string(),
Some("openai".to_string()),
Some("search".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(1),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn key() -> StoredProviderCatalogKey {
let auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","account_id":"account-search-1","is_fedramp":true}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-codex-search-1".to_string(),
"provider-codex-search-1".to_string(),
"oauth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:responses"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"codex-search-access-token",
)
.expect("access token should encrypt"),
Some(auth_config),
None,
Some(json!({"openai:search": 1})),
None,
Some(4_102_444_800),
None,
None,
)
.expect("key transport should build")
}
let seen_plans = Arc::new(Mutex::new(Vec::<serde_json::Value>::new()));
let seen_plans_clone = Arc::clone(&seen_plans);
let execution_runtime = Router::new().route(
@@ -316,7 +313,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
auth_snapshot(),
)]));
let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed({
let primary = candidate_row();
let primary = candidate_row("openai:search");
let mut backup = primary.clone();
backup.provider_id = "provider-codex-search-2".to_string();
backup.provider_name = "codex-backup".to_string();
@@ -338,7 +335,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
vec![primary, backup]
},
{
let primary = endpoint();
let primary = endpoint("openai:search");
let mut backup = primary.clone();
backup.id = "endpoint-codex-search-2".to_string();
backup.provider_id = "provider-codex-search-2".to_string();
@@ -598,3 +595,120 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_executes_codex_memories_with_responses_permission_and_native_json() {
run_search_sync_test(
"gateway_executes_codex_memories_with_responses_permission_and_native_json",
|| async {
let response_body = json!({"output":[{"trace_summary":"synthetic trace", "memory_summary":"synthetic memory"}],"future_response_field":{"enabled":true}});
let seen = Arc::new(Mutex::new(None));
let captured = Arc::clone(&seen);
let expected = response_body.clone();
let runtime = Router::new().route("/v1/execute/sync", any(move |request: Request| {
let captured = Arc::clone(&captured);
let response_body = expected.clone();
async move {
let bytes = to_bytes(request.into_body(), usize::MAX).await.expect("read plan");
let plan: serde_json::Value = serde_json::from_slice(&bytes).expect("parse plan");
let request_id = plan["request_id"].clone();
*captured.lock().expect("capture lock") = Some(plan);
let (status_code, response_body) = if request_id == json!("trace-memory-error") {
(400, json!({"error":{"type":"invalid_request_error","message":"synthetic invalid trace","code":"invalid_trace"},"future_error_field":{"enabled":true}}))
} else { (200, response_body) };
Json(json!({"request_id":request_id,"status_code":status_code,"headers":{"content-type":"application/json"},"body":{"json_body":response_body},"telemetry":{"elapsed_ms":1}}))
}
}));
let client_key = "sk-synthetic-memory-client";
let auth = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_key)),
auth_snapshot(),
)]));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row("openai:responses"),
]));
let mut memory_provider = provider();
memory_provider.config = Some(
json!({"codex":{"fingerprint_convergence_enabled":true},"failover_rules":{"stop_status_codes":[400]}}),
);
let catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![memory_provider],
vec![endpoint("openai:responses")],
vec![key()],
));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::default());
let data = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(auth,candidates,catalog,Arc::clone(&request_candidates),DEVELOPMENT_ENCRYPTION_KEY)
.with_system_config_values_for_tests([(crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string(),json!(true))]);
let (runtime_url, runtime_handle) = start_server(runtime).await;
let state = build_state_with_execution_runtime_override(runtime_url)
.with_data_state_for_tests(data);
let (url, gateway_handle) = start_server(build_router_with_state(state)).await;
let input = json!({"model":"gpt-5.6-sol-max","traces":[{"id":"synthetic-trace","metadata":{"source_path":"/synthetic/trace.json"},"items":[{"type":"message","role":"user","content":[]}]}],"reasoning":{"effort":"low"},"future_request_field":{"enabled":true}});
let response = reqwest::Client::new()
.post(format!("{url}/v1/memories/trace_summarize"))
.header(http::header::AUTHORIZATION, format!("Bearer {client_key}"))
.header(TRACE_ID_HEADER, "trace-memory-1")
.json(&input)
.send()
.await
.expect("send request");
let status = response.status();
let body: serde_json::Value = response.json().await.expect("read response");
assert_eq!(status, StatusCode::OK, "{body}");
assert_eq!(body, response_body);
let plan = seen
.lock()
.expect("capture lock")
.clone()
.expect("captured plan");
assert_eq!(
plan["url"],
"https://chatgpt.com/backend-api/codex/memories/trace_summarize"
);
assert_eq!(plan["stream"], false);
assert_eq!(plan["client_api_format"], "openai:responses");
assert_eq!(plan["provider_api_format"], "openai:responses");
assert_eq!(plan["headers"]["originator"], "codex_cli_rs");
assert_eq!(
plan["headers"]["authorization"],
"Bearer codex-search-access-token"
);
assert_eq!(plan["headers"]["accept"], "application/json");
assert!(plan["headers"]
.get("x-openai-internal-codex-responses-lite")
.is_none());
let mut expected_input = input;
expected_input["model"] = json!("gpt-5.6-sol");
expected_input["reasoning"]["effort"] = json!("max");
assert_eq!(plan["body"]["json_body"], expected_input);
let stored = request_candidates
.list_by_request_id("trace-memory-1")
.await
.expect("read candidates");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].status, RequestCandidateStatus::Success);
let error = reqwest::Client::new()
.post(format!("{url}/v1/memories/trace_summarize"))
.header(http::header::AUTHORIZATION, format!("Bearer {client_key}"))
.header(TRACE_ID_HEADER, "trace-memory-error")
.json(&expected_input)
.send()
.await
.expect("error response");
let error_status = error.status();
assert_eq!(error_status, StatusCode::BAD_REQUEST);
assert_eq!(
error.json::<serde_json::Value>().await.expect("error JSON"),
json!({"error":{"type":"invalid_request_error","message":"synthetic invalid trace","code":"invalid_trace"},"future_error_field":{"enabled":true}})
);
let error_candidates = request_candidates
.list_by_request_id("trace-memory-error")
.await
.expect("error candidate");
assert_eq!(error_candidates.len(), 1);
assert_eq!(error_candidates[0].status, RequestCandidateStatus::Failed);
gateway_handle.abort();
runtime_handle.abort();
},
);
}
@@ -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,

Some files were not shown because too many files have changed in this diff Show More