mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
Merge main into Claude Code dynamic client profile
This commit is contained in:
-1
@@ -1 +0,0 @@
|
||||
test binary
|
||||
@@ -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,
|
||||
|
||||
Vendored
+2
@@ -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
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
+3
-4
@@ -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()),
|
||||
};
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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)),
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user