feat: optimize usage body viewing and provider card layout

This commit is contained in:
elky
2026-09-08 02:49:06 +08:00
parent 7aa0c89244
commit 099b810a2f
36 changed files with 3113 additions and 739 deletions
@@ -1593,6 +1593,19 @@ impl GatewayDataState {
}
}
pub(crate) async fn read_request_usage_body_payload(
&self,
body_ref: &str,
) -> Result<
Option<aether_data_contracts::repository::usage::StoredUsageBodyPayload>,
DataLayerError,
> {
match &self.usage_reader {
Some(repository) => repository.read_body_payload(body_ref).await,
None => Ok(None),
}
}
pub(crate) async fn list_usage_audits(
&self,
query: &UsageAuditListQuery,
@@ -17,7 +17,10 @@ use aether_admin::observability::usage::{
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
admin_usage_provider_key_name, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageBodyField};
use aether_data_contracts::repository::usage::{
canonical_usage_body_ref_for, StoredRequestUsageAudit, StoredUsageBodyPayload,
UsageBodyCaptureState, UsageBodyField, MAX_DECOMPRESSED_USAGE_JSON_BYTES,
};
use axum::{
body::Body,
http,
@@ -28,9 +31,59 @@ use serde_json::{json, Value};
use std::collections::BTreeMap;
use tokio::try_join;
#[derive(Default)]
struct AdminUsageDetailBodyValue {
value: Option<Value>,
load_failed: bool,
error_code: Option<&'static str>,
}
impl AdminUsageDetailBodyValue {
fn resolved(
item: &StoredRequestUsageAudit,
field: UsageBodyField,
value: Option<Value>,
) -> Self {
let missing = value.is_none()
&& item
.body_capture_result(field, item.body_value(field))
.available;
Self {
value,
error_code: missing.then_some("missing"),
}
}
}
fn admin_usage_body_load_error_code(error: &GatewayError) -> &'static str {
if let GatewayError::Internal(message) = error {
if message.contains("decompressed usage json exceeds ")
|| message.contains("encoded usage json exceeds ")
{
return "too_large";
}
if message.contains("failed to decompress usage json:")
|| message.contains("failed to parse decompressed usage json:")
{
return "decode_failed";
}
}
"storage_unavailable"
}
async fn resolve_admin_usage_detail_field(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
field: UsageBodyField,
selected_field: Option<UsageBodyField>,
) -> AdminUsageDetailBodyValue {
if selected_field.is_some_and(|selected| selected != field) {
return AdminUsageDetailBodyValue::default();
}
if field == UsageBodyField::RequestBody {
resolve_admin_usage_detail_request_body(state, item).await
} else {
resolve_admin_usage_detail_body_value(state, item, field).await
}
}
async fn resolve_admin_usage_detail_request_body(
@@ -38,10 +91,7 @@ async fn resolve_admin_usage_detail_request_body(
item: &StoredRequestUsageAudit,
) -> AdminUsageDetailBodyValue {
match admin_usage_resolve_request_capture_body_for_item(state, item, None).await {
Ok(body) => AdminUsageDetailBodyValue {
value: body,
load_failed: false,
},
Ok(body) => AdminUsageDetailBodyValue::resolved(item, UsageBodyField::RequestBody, body),
Err(err) => {
tracing::warn!(
error = ?err,
@@ -52,7 +102,9 @@ async fn resolve_admin_usage_detail_request_body(
);
let value = admin_usage_resolve_request_capture_body(item, None);
AdminUsageDetailBodyValue {
load_failed: value.is_none(),
error_code: value
.is_none()
.then(|| admin_usage_body_load_error_code(&err)),
value,
}
}
@@ -66,10 +118,7 @@ async fn resolve_admin_usage_detail_body_value(
) -> AdminUsageDetailBodyValue {
let inline_body = item.body_value(field);
match admin_usage_resolve_body_value(state, item, inline_body, field).await {
Ok(body) => AdminUsageDetailBodyValue {
value: body,
load_failed: false,
},
Ok(body) => AdminUsageDetailBodyValue::resolved(item, field, body),
Err(err) => {
tracing::warn!(
error = ?err,
@@ -80,13 +129,139 @@ async fn resolve_admin_usage_detail_body_value(
);
let value = inline_body.cloned();
AdminUsageDetailBodyValue {
load_failed: value.is_none(),
error_code: value
.is_none()
.then(|| admin_usage_body_load_error_code(&err)),
value,
}
}
}
}
async fn read_admin_usage_raw_body(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
field: UsageBodyField,
) -> Result<Option<StoredUsageBodyPayload>, GatewayError> {
if matches!(
item.body_state(field),
Some(
UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable
| UsageBodyCaptureState::None
)
) {
return Ok(None);
}
let inline_body = item.body_value(field);
let prefer_inline = matches!(
item.body_state(field),
Some(UsageBodyCaptureState::Inline | UsageBodyCaptureState::Truncated)
) && inline_body.is_some();
if !prefer_inline {
if let Some(body_ref) = item
.body_ref(field)
.and_then(|reference| canonical_usage_body_ref_for(reference, &item.request_id, field))
{
if let Some(payload) = state.read_request_usage_body_payload(&body_ref).await? {
return Ok(Some(payload));
}
}
}
let fallback = inline_body.cloned().or_else(|| {
(field == UsageBodyField::RequestBody)
.then(|| admin_usage_resolve_request_capture_body(item, None))
.flatten()
});
fallback
.map(|value| {
serde_json::to_vec(&value)
.map(StoredUsageBodyPayload::Json)
.map_err(|error| GatewayError::Internal(error.to_string()))
})
.transpose()
}
async fn build_admin_usage_raw_body_response(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
field: UsageBodyField,
) -> Response<Body> {
let result = read_admin_usage_raw_body(state, item, field).await;
let mut response = match result {
Ok(Some(payload)) => admin_usage_raw_payload_response(payload),
Ok(None) => admin_usage_raw_body_error(http::StatusCode::NOT_FOUND, "missing"),
Err(error) => {
tracing::warn!(error = ?error, usage_id = %item.id, field = field.as_storage_field(), "failed to read admin usage raw body");
let code = admin_usage_body_load_error_code(&error);
admin_usage_raw_body_error(
if code == "too_large" {
http::StatusCode::PAYLOAD_TOO_LARGE
} else {
http::StatusCode::SERVICE_UNAVAILABLE
},
code,
)
}
};
let headers = response.headers_mut();
headers.insert(
http::header::CACHE_CONTROL,
http::HeaderValue::from_static("no-store, no-transform"),
);
headers.insert(
"x-content-type-options",
http::HeaderValue::from_static("nosniff"),
);
headers.insert(
"x-aether-body-field",
http::HeaderValue::from_static(field.as_storage_field()),
);
if let Ok(value) = http::HeaderValue::from_str(&item.id) {
headers.insert("x-aether-usage-id", value);
}
attach_admin_audit_response(
response,
"admin_usage_detail_viewed",
"view_usage_detail",
"usage_record",
&item.id,
)
}
fn admin_usage_raw_payload_response(payload: StoredUsageBodyPayload) -> Response<Body> {
let (encoding, bytes, limit) = match payload {
StoredUsageBodyPayload::Gzip(bytes) => (
"gzip",
bytes,
MAX_DECOMPRESSED_USAGE_JSON_BYTES + 1024 * 1024,
),
StoredUsageBodyPayload::Json(bytes) => ("json", bytes, MAX_DECOMPRESSED_USAGE_JSON_BYTES),
};
if bytes.len() > limit {
admin_usage_raw_body_error(http::StatusCode::PAYLOAD_TOO_LARGE, "too_large")
} else {
(
[
("content-type", "application/octet-stream"),
("content-encoding", "identity"),
("x-aether-body-encoding", encoding),
],
bytes,
)
.into_response()
}
}
fn admin_usage_raw_body_error(status: http::StatusCode, code: &'static str) -> Response<Body> {
(
status,
[("x-aether-body-error", code)],
Json(json!({ "body_load_error_code": code })),
)
.into_response()
}
pub(super) async fn maybe_build_local_admin_usage_detail_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -212,6 +387,25 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
"include_bodies",
true,
);
let body_field =
match request_context
.request_query_string
.as_deref()
.and_then(|query| {
url::form_urlencoded::parse(query.as_bytes())
.find(|(key, _)| key == "body_field")
.map(|(_, value)| value.into_owned())
}) {
Some(value) => {
match UsageBodyField::from_storage_field(value.trim()) {
Some(field) if include_bodies => Some(field),
_ => return Ok(Some(admin_usage_bad_request_response(
"body_field 必须是有效的正文字段,且 include_bodies 必须为 true",
))),
}
}
None => None,
};
let Some(item) = state.find_request_usage_by_id(&usage_id).await? else {
return Ok(Some(
@@ -223,6 +417,27 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
));
};
let body_format = request_context
.request_query_string
.as_deref()
.and_then(|query| {
url::form_urlencoded::parse(query.as_bytes())
.find(|(key, _)| key == "body_format")
.map(|(_, value)| value.into_owned())
});
if let Some(format) = body_format {
if format != "raw" || body_field.is_none() {
return Ok(Some(admin_usage_bad_request_response(
"body_format=raw 必须指定 body_field",
)));
}
if let Some(field) = body_field {
return Ok(Some(
build_admin_usage_raw_body_response(state, &item, field).await,
));
}
}
let user_ids = item.user_id.clone().into_iter().collect::<Vec<_>>();
let (users_by_id, provider_key_names, api_key_names): (
BTreeMap<String, aether_data::repository::users::StoredUserSummary>,
@@ -250,23 +465,32 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
}
}
let mut body_load_errors = serde_json::Map::new();
let request_body = if include_bodies {
let mut body_load_error_codes = serde_json::Map::new();
let mut request_body = if include_bodies {
let (request_body, provider_request_body, response_body, client_response_body) = tokio::join!(
resolve_admin_usage_detail_request_body(state, &item),
resolve_admin_usage_detail_body_value(
resolve_admin_usage_detail_field(
state,
&item,
UsageBodyField::RequestBody,
body_field
),
resolve_admin_usage_detail_field(
state,
&item,
UsageBodyField::ProviderRequestBody,
body_field,
),
resolve_admin_usage_detail_body_value(
resolve_admin_usage_detail_field(
state,
&item,
UsageBodyField::ResponseBody,
body_field,
),
resolve_admin_usage_detail_body_value(
resolve_admin_usage_detail_field(
state,
&item,
UsageBodyField::ClientResponseBody,
body_field,
),
);
for (field, resolved) in [
@@ -275,20 +499,25 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
(UsageBodyField::ResponseBody, &response_body),
(UsageBodyField::ClientResponseBody, &client_response_body),
] {
if resolved.load_failed {
if let Some(error_code) = resolved.error_code {
body_load_errors.insert(field.as_storage_field().to_string(), json!(true));
body_load_error_codes
.insert(field.as_storage_field().to_string(), json!(error_code));
}
}
detail_item.provider_request_body = provider_request_body.value;
detail_item.response_body = response_body.value;
detail_item.client_response_body = client_response_body.value;
if body_field.is_none_or(|field| field == UsageBodyField::ProviderRequestBody) {
detail_item.provider_request_body = provider_request_body.value;
}
if body_field.is_none_or(|field| field == UsageBodyField::ResponseBody) {
detail_item.response_body = response_body.value;
}
if body_field.is_none_or(|field| field == UsageBodyField::ClientResponseBody) {
detail_item.client_response_body = client_response_body.value;
}
request_body.value
} else {
None
};
if include_bodies {
// request_body 已通过 request capture 解析;其余 detached body 在上方并行加载。
}
let default_headers = admin_usage_curl_headers();
let mut payload = build_admin_usage_detail_payload(
&detail_item,
@@ -297,15 +526,33 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
state.has_auth_user_data_reader(),
state.has_auth_api_key_data_reader(),
provider_key_name.as_deref(),
include_bodies,
request_body,
include_bodies && body_field.is_none(),
if body_field.is_none() {
request_body.take()
} else {
None
},
&default_headers,
);
if let Some(field) = body_field {
payload[field.as_storage_field()] = match field {
UsageBodyField::RequestBody => request_body,
UsageBodyField::ProviderRequestBody => detail_item.provider_request_body.take(),
UsageBodyField::ResponseBody => detail_item.response_body.take(),
UsageBodyField::ClientResponseBody => detail_item.client_response_body.take(),
}
.unwrap_or(Value::Null);
}
payload["body_load_errors"] = if include_bodies && !body_load_errors.is_empty() {
Value::Object(body_load_errors)
} else {
Value::Null
};
payload["body_load_error_codes"] = if body_load_error_codes.is_empty() {
Value::Null
} else {
Value::Object(body_load_error_codes)
};
return Ok(Some(attach_admin_audit_response(
Json(payload).into_response(),
@@ -320,3 +567,61 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
Ok(None)
}
#[cfg(test)]
mod tests {
use super::admin_usage_body_load_error_code;
use crate::GatewayError;
#[tokio::test]
async fn admin_usage_raw_body_does_not_decode_or_reencode_stored_bytes() {
use super::{admin_usage_raw_payload_response, StoredUsageBodyPayload};
for (payload, encoding, expected) in [
(
StoredUsageBodyPayload::Gzip(vec![31, 139, 8, 0, 1]),
"gzip",
vec![31, 139, 8, 0, 1],
),
(
StoredUsageBodyPayload::Json(b"{ \"untouched\" : true }".to_vec()),
"json",
b"{ \"untouched\" : true }".to_vec(),
),
] {
let response = admin_usage_raw_payload_response(payload);
assert_eq!(response.headers()["content-encoding"], "identity");
assert_eq!(response.headers()["x-aether-body-encoding"], encoding);
let bytes = axum::body::to_bytes(response.into_body(), 1024)
.await
.unwrap();
assert_eq!(bytes.as_ref(), expected.as_slice());
}
}
#[test]
fn body_load_errors_expose_safe_codes_instead_of_internal_messages() {
for (message, expected) in [
(
"unexpected database value: decompressed usage json exceeds 67108864 bytes",
"too_large",
),
(
"failed to decompress usage json: invalid gzip header",
"decode_failed",
),
(
"failed to parse decompressed usage json: invalid JSON",
"decode_failed",
),
(
"postgres error: private connection details",
"storage_unavailable",
),
] {
assert_eq!(
admin_usage_body_load_error_code(&GatewayError::Internal(message.to_string())),
expected
);
}
}
}
@@ -234,6 +234,20 @@ impl<'a> AdminAppState<'a> {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn read_request_usage_body_payload(
&self,
body_ref: &str,
) -> Result<
Option<aether_data_contracts::repository::usage::StoredUsageBodyPayload>,
GatewayError,
> {
self.app
.data
.read_request_usage_body_payload(body_ref)
.await
.map_err(|error| GatewayError::Internal(error.to_string()))
}
pub(crate) async fn build_api_format_health_monitor_payload(
&self,
lookback_hours: u64,
@@ -211,9 +211,19 @@ fn apply_sensitive_route_cache_policy(
return;
}
let preserve_no_transform = headers
.get_all(http::header::CACHE_CONTROL)
.iter()
.filter_map(|value| value.to_str().ok())
.flat_map(|value| value.split(','))
.any(|directive| directive.trim().eq_ignore_ascii_case("no-transform"));
headers.insert(
http::header::CACHE_CONTROL,
HeaderValue::from_static("no-store"),
HeaderValue::from_static(if preserve_no_transform {
"no-store, no-transform"
} else {
"no-store"
}),
);
headers.insert(http::header::PRAGMA, HeaderValue::from_static("no-cache"));
}
@@ -354,6 +364,29 @@ mod tests {
);
}
#[test]
fn raw_body_no_transform_survives_sensitive_cache_policy() {
let mut headers = HeaderMap::new();
headers.append(
http::header::CACHE_CONTROL,
HeaderValue::from_static("public, max-age=3600"),
);
headers.append(
http::header::CACHE_CONTROL,
HeaderValue::from_static(" No-Transform "),
);
apply_sensitive_route_cache_policy(
&mut headers,
"/api/admin/usage/usage-1?body_format=raw",
None,
);
assert_eq!(
headers[http::header::CACHE_CONTROL],
"no-store, no-transform"
);
assert_eq!(headers[http::header::PRAGMA], "no-cache");
}
#[test]
fn authenticated_user_data_responses_are_never_cacheable() {
let mut headers = HeaderMap::new();
@@ -31,7 +31,11 @@ fn apply_frontdoor_cors_headers(
);
headers.insert(
http::header::ACCESS_CONTROL_EXPOSE_HEADERS,
HeaderValue::from_static("*"),
HeaderValue::from_static(if headers.contains_key("x-aether-body-field") {
"*, X-Aether-Body-Encoding, X-Aether-Body-Field, X-Aether-Body-Error, X-Aether-Usage-Id"
} else {
"*"
}),
);
if let Some(value) = requested_headers {
if let Ok(value) = HeaderValue::from_str(value) {
@@ -109,3 +113,36 @@ pub(crate) async fn frontdoor_cors_middleware(
);
response
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn raw_body_headers_are_explicitly_exposed_for_credentialed_requests() {
let cors =
FrontdoorCorsConfig::new(vec!["https://console.example".to_string()], true).unwrap();
let mut headers = http::HeaderMap::new();
headers.insert(
"x-aether-body-field",
HeaderValue::from_static("request_body"),
);
apply_frontdoor_cors_headers(&mut headers, &cors, "https://console.example", None);
let exposed = headers[http::header::ACCESS_CONTROL_EXPOSE_HEADERS]
.to_str()
.unwrap()
.to_ascii_lowercase();
for name in [
"x-aether-body-encoding",
"x-aether-body-field",
"x-aether-body-error",
"x-aether-usage-id",
] {
assert!(exposed.split(',').any(|header| header.trim() == name));
}
assert_eq!(
headers[http::header::ACCESS_CONTROL_ALLOW_CREDENTIALS],
"true"
);
}
}
@@ -2366,6 +2366,265 @@ async fn gateway_handles_admin_usage_detail_with_ref_backed_bodies() {
upstream_handle.abort();
}
fn sample_selective_body_usage() -> StoredRequestUsageAudit {
let mut usage = sample_usage_row(
"usage-selected-body",
"req-selected-body",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"completed",
120,
30,
0.3,
0.36,
DAY_1_UNIX_SECS,
);
usage.request_body = Some(json!({ "marker": "request_body" }));
usage.provider_request_body = Some(json!({ "marker": "provider_request_body" }));
usage.response_body = Some(json!({ "marker": "response_body" }));
usage.client_response_body = Some(json!({ "marker": "client_response_body" }));
usage
}
#[tokio::test]
async fn gateway_admin_usage_detail_returns_only_the_selected_body() {
let fields = [
"request_body",
"provider_request_body",
"response_body",
"client_response_body",
];
for detached in [false, true] {
let usage = sample_selective_body_usage();
let repository = if detached {
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage])
} else {
InMemoryUsageReadRepository::seed(vec![usage])
};
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
repository,
)));
for selected in fields {
let response = local_admin_usage_response(
&state, http::Method::GET,
&format!("/api/admin/usage/usage-selected-body?include_bodies=true&body_field={selected}"),
None,
).await;
assert_eq!(response.status(), StatusCode::OK);
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should collect");
let payload: serde_json::Value =
serde_json::from_slice(&bytes).expect("json should parse");
for field in fields {
assert_eq!(payload[format!("has_{field}")], true);
if field == selected {
assert_eq!(payload[field]["marker"], selected);
} else {
assert!(
payload[field].is_null(),
"unselected {field} must not be returned"
);
}
}
assert!(payload["body_load_errors"].is_null());
assert!(payload["body_load_error_codes"].is_null());
}
}
}
#[tokio::test]
async fn gateway_admin_usage_detail_validates_body_field_selection() {
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed(vec![sample_selective_body_usage()]),
)));
for query in [
"body_field=unknown",
"body_field=",
"include_bodies=false&body_field=request_body",
"body_format=raw",
"body_format=unknown&body_field=request_body",
"body_format=&body_field=request_body",
] {
let response = local_admin_usage_response(
&state,
http::Method::GET,
&format!("/api/admin/usage/usage-selected-body?{query}"),
None,
)
.await;
assert_eq!(
response.status(),
StatusCode::BAD_REQUEST,
"invalid query: {query}"
);
}
}
#[tokio::test]
async fn gateway_admin_usage_detail_raw_reads_only_selected_body_and_is_not_cacheable() {
for detached in [false, true] {
let repository = if detached {
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![
sample_selective_body_usage(),
])
} else {
InMemoryUsageReadRepository::seed(vec![sample_selective_body_usage()])
};
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::with_usage_reader_for_tests(Arc::new(repository)),
);
for field in [
"request_body",
"provider_request_body",
"response_body",
"client_response_body",
] {
let response = local_admin_usage_response(
&state,
http::Method::GET,
&format!("/api/admin/usage/usage-selected-body?body_field={field}&body_format=raw"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers()["cache-control"],
"no-store, no-transform"
);
assert_eq!(
response.headers()["x-aether-usage-id"],
"usage-selected-body"
);
assert_eq!(response.headers()["x-aether-body-field"], field);
assert_eq!(response.headers()["x-aether-body-encoding"], "json");
assert!(response.extensions().get::<AdminAuditEvent>().is_some());
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
json!({ "marker": field })
);
}
}
}
#[tokio::test]
async fn gateway_admin_usage_detail_raw_rejects_foreign_refs_and_disabled_capture() {
for body_state in [
UsageBodyCaptureState::Reference,
UsageBodyCaptureState::Disabled,
] {
let mut usage = sample_selective_body_usage();
usage.response_body = None;
usage.response_body_ref = Some("usage://request/foreign-request/response_body".to_string());
usage.response_body_state = Some(body_state);
let mut foreign = sample_selective_body_usage();
foreign.id = "foreign-usage".to_string();
foreign.request_id = "foreign-request".to_string();
foreign.response_body = Some(json!({ "secret": "must not leak" }));
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage, foreign]),
)),
);
let response = local_admin_usage_response(
&state,
http::Method::GET,
"/api/admin/usage/usage-selected-body?body_field=response_body&body_format=raw",
None,
)
.await;
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(response.headers()["x-aether-body-error"], "missing");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert!(!String::from_utf8_lossy(&bytes).contains("secret"));
}
}
#[tokio::test]
async fn gateway_admin_usage_detail_raw_preserves_authorization_and_binary_headers() {
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![
sample_selective_body_usage(),
]),
)),
);
let gateway =
build_router_with_state(state).layer(tower_http::compression::CompressionLayer::new());
let (url, server) = start_server(gateway).await;
let client = reqwest::Client::new();
let endpoint = format!(
"{url}/api/admin/usage/usage-selected-body?body_field=response_body&body_format=raw"
);
let unauthorized = client.get(&endpoint).send().await.unwrap();
assert!(matches!(
unauthorized.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
));
let response = admin_request(client.get(&endpoint))
.header("accept-encoding", "gzip")
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers()["content-encoding"], "identity");
assert_eq!(response.headers()["x-aether-body-encoding"], "json");
assert_eq!(
response.headers()["cache-control"],
"no-store, no-transform"
);
let bytes = response.bytes().await.unwrap();
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
json!({ "marker": "response_body" })
);
server.abort();
}
#[tokio::test]
async fn gateway_admin_usage_detail_isolates_missing_body_errors() {
let mut usage = sample_selective_body_usage();
usage.request_body = None;
usage.request_body_ref = Some("usage://request/req-selected-body/request_body".to_string());
usage.request_body_state = Some(UsageBodyCaptureState::Reference);
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage]),
)));
for selected in ["response_body", "request_body"] {
let response = local_admin_usage_response(
&state,
http::Method::GET,
&format!("/api/admin/usage/usage-selected-body?body_field={selected}"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should collect");
let payload: serde_json::Value = serde_json::from_slice(&bytes).expect("json should parse");
if selected == "request_body" {
assert_eq!(payload["body_load_errors"]["request_body"], true);
assert_eq!(payload["body_load_error_codes"]["request_body"], "missing");
assert!(payload["request_body"].is_null());
} else {
assert_eq!(payload["response_body"]["marker"], "response_body");
assert!(payload["body_load_errors"].is_null());
assert!(payload["body_load_error_codes"].is_null());
}
}
}
#[tokio::test]
async fn gateway_resolves_admin_usage_detail_when_inline_state_has_body_ref() {
let (_upstream_url, upstream_hits, upstream_handle) =