i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 18:07:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): enrich audit metadata and detail views

Browse Source
This commit is contained in:
elky
2026-07-17 19:20:16 +08:00
parent 0be380243b
commit 664c063a06
79 changed files with 6968 additions and 1294 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/usage/mod.rs
+389 -155
View File
@@ -49,7 +49,9 @@ use aether_data_contracts::repository::usage::{
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredUsageDailySummary,
UpsertUsageRecord, UsageAuditListQuery, UsageCounterFlushSummary, UsageCounterHealthSnapshot,
UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery, UsageReadRepository,
UsageWriteRepository,
UsageWriteRepository, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REQUESTED_REASONING_EFFORT_METADATA_KEY,
};
use aether_data_contracts::DataLayerError;
@@ -7994,8 +7996,14 @@ ORDER BY "usage".user_id ASC
http_audit_capture_mode,
routing_snapshot,
settlement_pricing_snapshot,
request_metadata_value,
request_metadata_json,
mut request_metadata_value,
mut request_metadata_json,
replace_client_request_body_facts,
replace_provider_request_body_facts,
clear_request_body,
clear_provider_request_body,
clear_response_body,
clear_client_response_body,
} = prepared;
Box::pin(async move {
lock_usage_request_id_in_tx(tx, &usage.request_id).await?;
@@ -8018,6 +8026,37 @@ ORDER BY "usage".user_id ASC
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
let capture_update_allowed = usage_capture_update_allowed(
previous_usage
.as_ref()
.map(|stored| (stored.status.as_str(), stored.billing_status.as_str())),
usage.status.as_str(),
);
let replace_terminal_snapshots =
matches!(usage.status.as_str(), "completed" | "failed" | "cancelled");
if capture_update_allowed
&& request_metadata_value.is_none()
&& (replace_terminal_snapshots
|| replace_client_request_body_facts
|| replace_provider_request_body_facts)
{
let previous_metadata = previous_usage
.as_ref()
.and_then(|stored| stored.request_metadata.as_ref());
request_metadata_value = Some(if replace_terminal_snapshots {
retain_previous_request_audit_metadata(
previous_metadata,
!replace_client_request_body_facts,
)
} else {
clear_previous_request_body_facts(
previous_metadata,
replace_client_request_body_facts,
replace_provider_request_body_facts,
)
});
request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
}
let _row = sqlx::query(UPSERT_SQL)
.bind(Uuid::new_v4().to_string())
.bind(&usage.request_id)
@@ -8090,72 +8129,89 @@ ORDER BY "usage".user_id ASC
usage.updated_at_unix_secs
))
})?)
.bind(request_body_storage.has_detached_blob())
.bind(provider_request_body_storage.has_detached_blob())
.bind(response_body_storage.has_detached_blob())
.bind(client_response_body_storage.has_detached_blob())
.bind(request_body_storage.has_detached_blob() || clear_request_body)
.bind(
provider_request_body_storage.has_detached_blob()
|| clear_provider_request_body,
)
.bind(response_body_storage.has_detached_blob() || clear_response_body)
.bind(
client_response_body_storage.has_detached_blob()
|| clear_client_response_body,
)
.bind(capture_update_allowed)
.fetch_one(&mut **tx)
.await
.map_postgres_err()?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::RequestBody,
usage.request_body.as_ref(),
&request_body_storage,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ProviderRequestBody,
usage.provider_request_body.as_ref(),
&provider_request_body_storage,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ResponseBody,
usage.response_body.as_ref(),
&response_body_storage,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ClientResponseBody,
usage.client_response_body.as_ref(),
&client_response_body_storage,
)
.await?;
let http_audit_headers = UsageHttpAuditHeaders {
request_headers_json: request_headers_json.as_deref(),
provider_request_headers_json: provider_request_headers_json.as_deref(),
response_headers_json: response_headers_json.as_deref(),
client_response_headers_json: client_response_headers_json.as_deref(),
};
sync_usage_http_audit_storage(
&mut **tx,
&usage.request_id,
&http_audit_headers,
&http_audit_refs,
&http_audit_states,
http_audit_capture_mode,
)
.await?;
sync_usage_routing_snapshot_storage(
&mut **tx,
&usage.request_id,
&routing_snapshot,
)
.await?;
sync_usage_settlement_pricing_snapshot_storage(
&mut **tx,
&usage.request_id,
&settlement_pricing_snapshot,
)
.await?;
if capture_update_allowed {
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::RequestBody,
usage.request_body.as_ref(),
&request_body_storage,
clear_request_body,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ProviderRequestBody,
usage.provider_request_body.as_ref(),
&provider_request_body_storage,
clear_provider_request_body,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ResponseBody,
usage.response_body.as_ref(),
&response_body_storage,
clear_response_body,
)
.await?;
sync_usage_body_blob_storage(
&mut **tx,
&usage.request_id,
UsageBodyField::ClientResponseBody,
usage.client_response_body.as_ref(),
&client_response_body_storage,
clear_client_response_body,
)
.await?;
let http_audit_headers = UsageHttpAuditHeaders {
request_headers_json: request_headers_json.as_deref(),
provider_request_headers_json: provider_request_headers_json.as_deref(),
response_headers_json: response_headers_json.as_deref(),
client_response_headers_json: client_response_headers_json.as_deref(),
};
sync_usage_http_audit_storage(
&mut **tx,
&usage.request_id,
&http_audit_headers,
&http_audit_refs,
&http_audit_states,
http_audit_capture_mode,
)
.await?;
}
if capture_update_allowed {
sync_usage_routing_snapshot_storage(
&mut **tx,
&usage.request_id,
&routing_snapshot,
replace_terminal_snapshots,
)
.await?;
sync_usage_settlement_pricing_snapshot_storage(
&mut **tx,
&usage.request_id,
&settlement_pricing_snapshot,
replace_terminal_snapshots,
)
.await?;
}
let mut stored = find_usage_by_request_id_in_tx(tx, &usage.request_id)
.await?
@@ -8165,71 +8221,79 @@ ORDER BY "usage".user_id ASC
usage.request_id
))
})?;
if request_body_storage.has_detached_blob() {
stored.request_body = usage.request_body.clone();
if capture_update_allowed {
if request_body_storage.has_detached_blob() {
stored.request_body = usage.request_body.clone();
}
stored.request_headers = usage.request_headers.clone();
stored.provider_request_headers = usage.provider_request_headers.clone();
if provider_request_body_storage.has_detached_blob() {
stored.provider_request_body = usage.provider_request_body.clone();
}
stored.response_headers = usage.response_headers.clone();
if response_body_storage.has_detached_blob() {
stored.response_body = usage.response_body.clone();
}
stored.client_response_headers = usage.client_response_headers.clone();
if client_response_body_storage.has_detached_blob() {
stored.client_response_body = usage.client_response_body.clone();
}
stored.request_body_ref = if clear_request_body {
None
} else {
resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
http_audit_refs.request_body_ref.as_deref(),
)
};
stored.provider_request_body_ref = if clear_provider_request_body {
None
} else {
resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
http_audit_refs.provider_request_body_ref.as_deref(),
)
};
stored.response_body_ref = if clear_response_body {
None
} else {
resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
http_audit_refs.response_body_ref.as_deref(),
)
};
stored.client_response_body_ref = if clear_client_response_body {
None
} else {
resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
http_audit_refs.client_response_body_ref.as_deref(),
)
};
stored.request_body_state =
usage.request_body_state.or(stored.request_body_state);
stored.provider_request_body_state = usage
.provider_request_body_state
.or(stored.provider_request_body_state);
stored.response_body_state =
usage.response_body_state.or(stored.response_body_state);
stored.client_response_body_state = usage
.client_response_body_state
.or(stored.client_response_body_state);
stored.request_metadata = request_metadata_value;
}
stored.request_headers = usage.request_headers.clone();
stored.provider_request_headers = usage.provider_request_headers.clone();
if provider_request_body_storage.has_detached_blob() {
stored.provider_request_body = usage.provider_request_body.clone();
}
stored.response_headers = usage.response_headers.clone();
if response_body_storage.has_detached_blob() {
stored.response_body = usage.response_body.clone();
}
stored.client_response_headers = usage.client_response_headers.clone();
if client_response_body_storage.has_detached_blob() {
stored.client_response_body = usage.client_response_body.clone();
}
stored.request_body_ref = resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
http_audit_refs.request_body_ref.as_deref(),
);
stored.provider_request_body_ref = resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
http_audit_refs.provider_request_body_ref.as_deref(),
);
stored.response_body_ref = resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
http_audit_refs.response_body_ref.as_deref(),
);
stored.client_response_body_ref = resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
http_audit_refs.client_response_body_ref.as_deref(),
);
stored.request_body_state =
usage.request_body_state.or(stored.request_body_state);
stored.provider_request_body_state = usage
.provider_request_body_state
.or(stored.provider_request_body_state);
stored.response_body_state =
usage.response_body_state.or(stored.response_body_state);
stored.client_response_body_state = usage
.client_response_body_state
.or(stored.client_response_body_state);
stored.candidate_id = routing_snapshot.candidate_id.clone();
stored.candidate_index = routing_snapshot.candidate_index;
stored.key_name = routing_snapshot.key_name.clone();
stored.planner_kind = routing_snapshot.planner_kind.clone();
stored.route_family = routing_snapshot.route_family.clone();
stored.route_kind = routing_snapshot.route_kind.clone();
stored.execution_path = routing_snapshot.execution_path.clone();
stored.local_execution_runtime_miss_reason =
routing_snapshot.local_execution_runtime_miss_reason.clone();
stored.output_price_per_1m = settlement_pricing_snapshot.output_price_per_1m;
stored.request_metadata = request_metadata_value;
let before_api_key_contribution =
previous_usage.as_ref().and_then(api_key_usage_contribution);
@@ -10415,6 +10479,112 @@ struct PreparedUsageUpsert {
settlement_pricing_snapshot: UsageSettlementPricingSnapshot,
request_metadata_value: Option<Value>,
request_metadata_json: Option<String>,
replace_client_request_body_facts: bool,
replace_provider_request_body_facts: bool,
clear_request_body: bool,
clear_provider_request_body: bool,
clear_response_body: bool,
clear_client_response_body: bool,
}
fn request_body_capture_replaces_derived_facts(
request_body: Option<&Value>,
request_body_state: Option<UsageBodyCaptureState>,
) -> bool {
// A typed capture state belongs to the incoming request snapshot. Metadata derived before a
// body was externalized, truncated, disabled, or found unavailable is authoritative when
// present; its absence must clear facts from an older candidate instead of falling through to
// PostgreSQL's sparse-upsert COALESCE behavior.
if request_body_state.is_some() {
return true;
}
let Some(request_body) = request_body else {
return false;
};
!request_body.as_object().is_some_and(|body| {
body.get("truncated").and_then(Value::as_bool) == Some(true)
&& body.get("reason").and_then(Value::as_str) == Some("body_capture_limit_exceeded")
})
}
fn clear_previous_request_body_facts(
previous_metadata: Option<&Value>,
clear_client_request_body_facts: bool,
clear_provider_request_body_facts: bool,
) -> Value {
let mut metadata = previous_metadata
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
if clear_client_request_body_facts {
metadata.remove(REQUESTED_REASONING_EFFORT_METADATA_KEY);
metadata.remove("request_body_ref");
}
if clear_provider_request_body_facts {
metadata.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY);
metadata.remove(PROVIDER_SERVICE_TIER_METADATA_KEY);
metadata.remove(PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY);
metadata.remove("provider_request_body_ref");
}
// Keep an explicit empty object as a tombstone. Binding SQL NULL here would make the upsert's
// COALESCE retain the previous candidate's request-derived facts.
Value::Object(metadata)
}
fn retain_previous_request_audit_metadata(
previous_metadata: Option<&Value>,
preserve_client_request_body_facts: bool,
) -> Value {
let Some(previous_metadata) = previous_metadata.and_then(Value::as_object) else {
return Value::Object(Map::new());
};
let mut retained = Map::new();
for key in [
"trace_id",
"client_ip",
"user_agent",
"client_family",
"client_requested_stream",
"client_session_affinity",
"api_key_is_standalone",
"request_path",
"request_query_string",
"request_path_and_query",
] {
if let Some(value) = previous_metadata.get(key) {
retained.insert(key.to_string(), value.clone());
}
}
if preserve_client_request_body_facts {
for key in [REQUESTED_REASONING_EFFORT_METADATA_KEY, "request_body_ref"] {
if let Some(value) = previous_metadata.get(key) {
retained.insert(key.to_string(), value.clone());
}
}
}
Value::Object(retained)
}
fn usage_capture_update_allowed(
previous_usage: Option<(&str, &str)>,
incoming_status: &str,
) -> bool {
let Some((previous_status, previous_billing_status)) = previous_usage else {
return true;
};
if previous_billing_status != "pending" {
return false;
}
let previous_is_terminal = matches!(previous_status, "completed" | "failed" | "cancelled");
let incoming_is_non_terminal = matches!(incoming_status, "pending" | "streaming");
if previous_is_terminal && incoming_is_non_terminal {
return false;
}
!(previous_status == "streaming" && incoming_status == "pending")
}
fn prepare_usage_body_storage(value: Option<&Value>) -> Result<UsageBodyStorage, DataLayerError> {
@@ -10485,40 +10655,78 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
fn prepare_usage_upsert_context(
usage: &UpsertUsageRecord,
) -> Result<PreparedUsageUpsert, DataLayerError> {
let replace_client_request_body_facts = request_body_capture_replaces_derived_facts(
usage.request_body.as_ref(),
usage.request_body_state,
);
let replace_provider_request_body_facts = request_body_capture_replaces_derived_facts(
usage.provider_request_body.as_ref(),
usage.provider_request_body_state,
);
let clear_request_body = usage.request_body_state == Some(UsageBodyCaptureState::None);
let clear_provider_request_body =
usage.provider_request_body_state == Some(UsageBodyCaptureState::None);
let clear_response_body = usage.response_body_state == Some(UsageBodyCaptureState::None);
let clear_client_response_body =
usage.client_response_body_state == Some(UsageBodyCaptureState::None);
// A typed `none` marker wins over residual values left on a reused event by an earlier
// candidate. Do not serialize those values or recreate their detached references.
let request_body = (!clear_request_body)
.then_some(usage.request_body.as_ref())
.flatten();
let provider_request_body = (!clear_provider_request_body)
.then_some(usage.provider_request_body.as_ref())
.flatten();
let response_body = (!clear_response_body)
.then_some(usage.response_body.as_ref())
.flatten();
let client_response_body = (!clear_client_response_body)
.then_some(usage.client_response_body.as_ref())
.flatten();
let request_body_ref = (!clear_request_body)
.then_some(usage.request_body_ref.as_deref())
.flatten();
let provider_request_body_ref = (!clear_provider_request_body)
.then_some(usage.provider_request_body_ref.as_deref())
.flatten();
let response_body_ref = (!clear_response_body)
.then_some(usage.response_body_ref.as_deref())
.flatten();
let client_response_body_ref = (!clear_client_response_body)
.then_some(usage.client_response_body_ref.as_deref())
.flatten();
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage = prepare_usage_body_storage(usage.request_body.as_ref())?;
let request_body_storage = prepare_usage_body_storage(request_body)?;
let provider_request_headers_json = json_bind_text(usage.provider_request_headers.as_ref())?;
let provider_request_body_storage =
prepare_usage_body_storage(usage.provider_request_body.as_ref())?;
let provider_request_body_storage = prepare_usage_body_storage(provider_request_body)?;
let response_headers_json = json_bind_text(usage.response_headers.as_ref())?;
let response_body_storage = prepare_usage_body_storage(usage.response_body.as_ref())?;
let response_body_storage = prepare_usage_body_storage(response_body)?;
let client_response_headers_json = json_bind_text(usage.client_response_headers.as_ref())?;
let client_response_body_storage =
prepare_usage_body_storage(usage.client_response_body.as_ref())?;
let client_response_body_storage = prepare_usage_body_storage(client_response_body)?;
let http_audit_refs = UsageHttpAuditRefs {
request_body_ref: resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
request_body_ref,
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
None,
),
provider_request_body_ref: resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
provider_request_body_ref,
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
None,
),
response_body_ref: resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
response_body_ref,
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
None,
),
client_response_body_ref: resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
client_response_body_ref,
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
@@ -10547,42 +10755,49 @@ fn prepare_usage_upsert_context(
http_audit_refs.client_response_body_ref.as_deref(),
),
};
let request_metadata_value = prepare_request_metadata_for_body_storage(
let mut request_metadata_value = prepare_request_metadata_for_body_storage(
usage.request_metadata.clone(),
[
(
UsageBodyField::RequestBody,
&request_body_storage,
usage.request_body.as_ref(),
usage.request_body_ref.as_deref(),
request_body,
request_body_ref,
),
(
UsageBodyField::ProviderRequestBody,
&provider_request_body_storage,
usage.provider_request_body.as_ref(),
usage.provider_request_body_ref.as_deref(),
provider_request_body,
provider_request_body_ref,
),
(
UsageBodyField::ResponseBody,
&response_body_storage,
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
response_body,
response_body_ref,
),
(
UsageBodyField::ClientResponseBody,
&client_response_body_storage,
usage.client_response_body.as_ref(),
usage.client_response_body_ref.as_deref(),
client_response_body,
client_response_body_ref,
),
],
);
if request_metadata_value.is_some() && (clear_request_body || clear_provider_request_body) {
request_metadata_value = Some(clear_previous_request_body_facts(
request_metadata_value.as_ref(),
clear_request_body,
clear_provider_request_body,
));
}
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
usage.request_body.as_ref(),
usage.provider_request_body.as_ref(),
usage.response_body.as_ref(),
usage.client_response_body.as_ref(),
request_body,
provider_request_body,
response_body,
client_response_body,
],
);
let routing_snapshot =
@@ -10607,6 +10822,12 @@ fn prepare_usage_upsert_context(
settlement_pricing_snapshot,
request_metadata_value,
request_metadata_json,
replace_client_request_body_facts,
replace_provider_request_body_facts,
clear_request_body,
clear_provider_request_body,
clear_response_body,
clear_client_response_body,
})
}
@@ -11442,11 +11663,20 @@ async fn sync_usage_body_blob_storage<'e, E>(
field: UsageBodyField,
value: Option<&Value>,
storage: &UsageBodyStorage,
clear_existing: bool,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
let body_ref = usage_body_ref(request_id, field);
if clear_existing {
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
return Ok(());
}
if let Some(payload_gzip) = storage.detached_blob_bytes.as_ref() {
sqlx::query(UPSERT_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
@@ -11523,11 +11753,12 @@ async fn sync_usage_routing_snapshot_storage<'e, E>(
executor: E,
request_id: &str,
snapshot: &UsageRoutingSnapshot,
replace_existing: bool,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if !snapshot.any_present() {
if !snapshot.any_present() && !replace_existing {
return Ok(());
}
@@ -11545,6 +11776,7 @@ where
.bind(snapshot.selected_endpoint_id.as_deref())
.bind(snapshot.selected_provider_api_key_id.as_deref())
.bind(snapshot.has_format_conversion)
.bind(replace_existing)
.execute(executor)
.await
.map_postgres_err()?;
@@ -11556,11 +11788,12 @@ async fn sync_usage_settlement_pricing_snapshot_storage<'e, E>(
executor: E,
request_id: &str,
snapshot: &UsageSettlementPricingSnapshot,
replace_existing: bool,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if !snapshot.any_present() {
if !snapshot.any_present() && !replace_existing {
return Ok(());
}
@@ -11594,6 +11827,7 @@ where
.bind(snapshot.cache_creation_price_per_1m)
.bind(snapshot.cache_read_price_per_1m)
.bind(snapshot.price_per_request)
.bind(replace_existing)
.execute(executor)
.await
.map_postgres_err()?;
crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql
+3
View File
@@ -179,6 +179,7 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'requested_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL
@@ -193,6 +194,8 @@ SELECT
NULLIF(BTRIM("usage".request_metadata->>'request_path'), ''),
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'requested_reasoning_effort',
NULLIF(BTRIM("usage".request_metadata->>'requested_reasoning_effort'), ''),
'provider_reasoning_effort',
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
'provider_service_tier',
crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql
+3
View File
@@ -179,6 +179,7 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'requested_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL
@@ -193,6 +194,8 @@ SELECT
NULLIF(BTRIM("usage".request_metadata->>'request_path'), ''),
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'requested_reasoning_effort',
NULLIF(BTRIM("usage".request_metadata->>'requested_reasoning_effort'), ''),
'provider_reasoning_effort',
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
'provider_service_tier',
crates/aether-data/adapters/postgres/src/usage/queries/upsert_sql.sql
+25 -25
View File
@@ -157,22 +157,22 @@ DO UPDATE SET
api_key_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_id, "usage".api_key_id) ELSE "usage".api_key_id END,
username = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.username, "usage".username) ELSE "usage".username END,
api_key_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_name, "usage".api_key_name) ELSE "usage".api_key_name END,
provider_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_name, "usage".provider_name) ELSE "usage".provider_name END,
model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.model, "usage".model) ELSE "usage".model END,
target_model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.target_model, "usage".target_model) ELSE "usage".target_model END,
provider_id = CASE WHEN "usage".billing_status = 'pending' OR ("usage".provider_id IS NULL AND ("usage".provider_endpoint_id IS NULL OR "usage".provider_endpoint_id = EXCLUDED.provider_endpoint_id) AND ("usage".provider_api_key_id IS NULL OR "usage".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_id, "usage".provider_id) ELSE "usage".provider_id END,
provider_endpoint_id = CASE WHEN "usage".billing_status = 'pending' OR ("usage".provider_endpoint_id IS NULL AND ("usage".provider_id IS NULL OR "usage".provider_id = EXCLUDED.provider_id) AND ("usage".provider_api_key_id IS NULL OR "usage".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_endpoint_id, "usage".provider_endpoint_id) ELSE "usage".provider_endpoint_id END,
provider_api_key_id = CASE WHEN "usage".billing_status = 'pending' OR ("usage".provider_api_key_id IS NULL AND ("usage".provider_id IS NULL OR "usage".provider_id = EXCLUDED.provider_id) AND ("usage".provider_endpoint_id IS NULL OR "usage".provider_endpoint_id = EXCLUDED.provider_endpoint_id)) THEN COALESCE(EXCLUDED.provider_api_key_id, "usage".provider_api_key_id) ELSE "usage".provider_api_key_id END,
request_type = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_type, "usage".request_type) ELSE "usage".request_type END,
api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_format, "usage".api_format) ELSE "usage".api_format END,
api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_family, "usage".api_family) ELSE "usage".api_family END,
endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_kind, "usage".endpoint_kind) ELSE "usage".endpoint_kind END,
endpoint_api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_api_format, "usage".endpoint_api_format) ELSE "usage".endpoint_api_format END,
provider_api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_api_family, "usage".provider_api_family) ELSE "usage".provider_api_family END,
provider_endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_kind, "usage".provider_endpoint_kind) ELSE "usage".provider_endpoint_kind END,
has_format_conversion = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.has_format_conversion, "usage".has_format_conversion) ELSE "usage".has_format_conversion END,
is_stream = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.is_stream, "usage".is_stream) ELSE "usage".is_stream END,
upstream_is_stream = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.upstream_is_stream, "usage".upstream_is_stream, "usage".is_stream, false) ELSE "usage".upstream_is_stream END,
provider_name = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.provider_name, "usage".provider_name) ELSE "usage".provider_name END,
model = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.model, "usage".model) ELSE "usage".model END,
target_model = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE WHEN EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN EXCLUDED.target_model ELSE COALESCE(EXCLUDED.target_model, "usage".target_model) END ELSE "usage".target_model END,
provider_id = CASE WHEN ("usage".billing_status = 'pending' AND $61) OR ("usage".billing_status <> 'pending' AND "usage".provider_id IS NULL AND ("usage".provider_endpoint_id IS NULL OR "usage".provider_endpoint_id = EXCLUDED.provider_endpoint_id) AND ("usage".provider_api_key_id IS NULL OR "usage".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_id, "usage".provider_id) ELSE "usage".provider_id END,
provider_endpoint_id = CASE WHEN ("usage".billing_status = 'pending' AND $61) OR ("usage".billing_status <> 'pending' AND "usage".provider_endpoint_id IS NULL AND ("usage".provider_id IS NULL OR "usage".provider_id = EXCLUDED.provider_id) AND ("usage".provider_api_key_id IS NULL OR "usage".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_endpoint_id, "usage".provider_endpoint_id) ELSE "usage".provider_endpoint_id END,
provider_api_key_id = CASE WHEN ("usage".billing_status = 'pending' AND $61) OR ("usage".billing_status <> 'pending' AND "usage".provider_api_key_id IS NULL AND ("usage".provider_id IS NULL OR "usage".provider_id = EXCLUDED.provider_id) AND ("usage".provider_endpoint_id IS NULL OR "usage".provider_endpoint_id = EXCLUDED.provider_endpoint_id)) THEN COALESCE(EXCLUDED.provider_api_key_id, "usage".provider_api_key_id) ELSE "usage".provider_api_key_id END,
request_type = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.request_type, "usage".request_type) ELSE "usage".request_type END,
api_format = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.api_format, "usage".api_format) ELSE "usage".api_format END,
api_family = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.api_family, "usage".api_family) ELSE "usage".api_family END,
endpoint_kind = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.endpoint_kind, "usage".endpoint_kind) ELSE "usage".endpoint_kind END,
endpoint_api_format = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.endpoint_api_format, "usage".endpoint_api_format) ELSE "usage".endpoint_api_format END,
provider_api_family = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.provider_api_family, "usage".provider_api_family) ELSE "usage".provider_api_family END,
provider_endpoint_kind = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.provider_endpoint_kind, "usage".provider_endpoint_kind) ELSE "usage".provider_endpoint_kind END,
has_format_conversion = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.has_format_conversion, "usage".has_format_conversion) ELSE "usage".has_format_conversion END,
is_stream = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.is_stream, "usage".is_stream) ELSE "usage".is_stream END,
upstream_is_stream = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.upstream_is_stream, "usage".upstream_is_stream, "usage".is_stream, false) ELSE "usage".upstream_is_stream END,
input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_tokens, EXCLUDED.input_tokens) ELSE "usage".input_tokens END,
output_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".output_tokens, EXCLUDED.output_tokens) ELSE "usage".output_tokens END,
total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".total_tokens, EXCLUDED.total_tokens) ELSE "usage".total_tokens END,
@@ -223,42 +223,42 @@ DO UPDATE SET
END ELSE "usage".status END,
billing_status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.billing_status ELSE "usage".billing_status END,
request_headers = NULL,
request_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
request_body = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.request_body_compressed IS NOT NULL OR $57 THEN NULL
ELSE COALESCE(EXCLUDED.request_body, "usage".request_body)
END ELSE "usage".request_body END,
request_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
request_body_compressed = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.request_body IS NOT NULL OR $57 THEN NULL
ELSE COALESCE(EXCLUDED.request_body_compressed, "usage".request_body_compressed)
END ELSE "usage".request_body_compressed END,
provider_request_headers = NULL,
provider_request_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
provider_request_body = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.provider_request_body_compressed IS NOT NULL OR $58 THEN NULL
ELSE COALESCE(EXCLUDED.provider_request_body, "usage".provider_request_body)
END ELSE "usage".provider_request_body END,
provider_request_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
provider_request_body_compressed = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.provider_request_body IS NOT NULL OR $58 THEN NULL
ELSE COALESCE(EXCLUDED.provider_request_body_compressed, "usage".provider_request_body_compressed)
END ELSE "usage".provider_request_body_compressed END,
response_headers = NULL,
response_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
response_body = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.response_body_compressed IS NOT NULL OR $59 THEN NULL
ELSE COALESCE(EXCLUDED.response_body, "usage".response_body)
END ELSE "usage".response_body END,
response_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
response_body_compressed = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.response_body IS NOT NULL OR $59 THEN NULL
ELSE COALESCE(EXCLUDED.response_body_compressed, "usage".response_body_compressed)
END ELSE "usage".response_body_compressed END,
client_response_headers = NULL,
client_response_body = CASE WHEN "usage".billing_status = 'pending' THEN CASE
client_response_body = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $60 THEN NULL
ELSE COALESCE(EXCLUDED.client_response_body, "usage".client_response_body)
END ELSE "usage".client_response_body END,
client_response_body_compressed = CASE WHEN "usage".billing_status = 'pending' THEN CASE
client_response_body_compressed = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN CASE
WHEN EXCLUDED.client_response_body IS NOT NULL OR $60 THEN NULL
ELSE COALESCE(EXCLUDED.client_response_body_compressed, "usage".client_response_body_compressed)
END ELSE "usage".client_response_body_compressed END,
request_metadata = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_metadata, "usage".request_metadata) ELSE "usage".request_metadata END,
request_metadata = CASE WHEN "usage".billing_status = 'pending' AND $61 THEN COALESCE(EXCLUDED.request_metadata, "usage".request_metadata) ELSE "usage".request_metadata END,
finalized_at = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.finalized_at, "usage".finalized_at) ELSE "usage".finalized_at END,
updated_at_unix_secs = CASE WHEN "usage".billing_status = 'pending' THEN
GREATEST(
crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_http_audit_sql.sql
+34 -15
View File
@@ -41,16 +41,25 @@ DO UPDATE SET
EXCLUDED.client_response_headers,
usage_http_audits.client_response_headers
),
request_body_ref = COALESCE(EXCLUDED.request_body_ref, usage_http_audits.request_body_ref),
provider_request_body_ref = COALESCE(
EXCLUDED.provider_request_body_ref,
usage_http_audits.provider_request_body_ref
),
response_body_ref = COALESCE(EXCLUDED.response_body_ref, usage_http_audits.response_body_ref),
client_response_body_ref = COALESCE(
EXCLUDED.client_response_body_ref,
usage_http_audits.client_response_body_ref
),
request_body_ref = CASE
WHEN EXCLUDED.request_body_state = 'none' THEN NULL
ELSE COALESCE(EXCLUDED.request_body_ref, usage_http_audits.request_body_ref)
END,
provider_request_body_ref = CASE
WHEN EXCLUDED.provider_request_body_state = 'none' THEN NULL
ELSE COALESCE(EXCLUDED.provider_request_body_ref, usage_http_audits.provider_request_body_ref)
END,
response_body_ref = CASE
WHEN EXCLUDED.response_body_state = 'none' THEN NULL
ELSE COALESCE(EXCLUDED.response_body_ref, usage_http_audits.response_body_ref)
END,
client_response_body_ref = CASE
WHEN EXCLUDED.client_response_body_state = 'none' THEN NULL
ELSE COALESCE(
EXCLUDED.client_response_body_ref,
usage_http_audits.client_response_body_ref
)
END,
request_body_state = COALESCE(
EXCLUDED.request_body_state,
usage_http_audits.request_body_state
@@ -67,9 +76,19 @@ DO UPDATE SET
EXCLUDED.client_response_body_state,
usage_http_audits.client_response_body_state
),
body_capture_mode = COALESCE(
NULLIF(EXCLUDED.body_capture_mode, 'none'),
usage_http_audits.body_capture_mode,
'none'
),
body_capture_mode = CASE
WHEN EXCLUDED.body_capture_mode = 'none'
AND (
EXCLUDED.request_body_state = 'none'
OR EXCLUDED.provider_request_body_state = 'none'
OR EXCLUDED.response_body_state = 'none'
OR EXCLUDED.client_response_body_state = 'none'
)
THEN 'none'
ELSE COALESCE(
NULLIF(EXCLUDED.body_capture_mode, 'none'),
usage_http_audits.body_capture_mode,
'none'
)
END,
updated_at = NOW()
crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_routing_snapshot_sql.sql
+12 -30
View File
@@ -29,34 +29,16 @@ INSERT INTO usage_routing_snapshots (
)
ON CONFLICT (request_id)
DO UPDATE SET
candidate_id = COALESCE(EXCLUDED.candidate_id, usage_routing_snapshots.candidate_id),
candidate_index = COALESCE(
EXCLUDED.candidate_index,
usage_routing_snapshots.candidate_index
),
key_name = COALESCE(EXCLUDED.key_name, usage_routing_snapshots.key_name),
planner_kind = COALESCE(EXCLUDED.planner_kind, usage_routing_snapshots.planner_kind),
route_family = COALESCE(EXCLUDED.route_family, usage_routing_snapshots.route_family),
route_kind = COALESCE(EXCLUDED.route_kind, usage_routing_snapshots.route_kind),
execution_path = COALESCE(EXCLUDED.execution_path, usage_routing_snapshots.execution_path),
local_execution_runtime_miss_reason = COALESCE(
EXCLUDED.local_execution_runtime_miss_reason,
usage_routing_snapshots.local_execution_runtime_miss_reason
),
selected_provider_id = COALESCE(
EXCLUDED.selected_provider_id,
usage_routing_snapshots.selected_provider_id
),
selected_endpoint_id = COALESCE(
EXCLUDED.selected_endpoint_id,
usage_routing_snapshots.selected_endpoint_id
),
selected_provider_api_key_id = COALESCE(
EXCLUDED.selected_provider_api_key_id,
usage_routing_snapshots.selected_provider_api_key_id
),
has_format_conversion = COALESCE(
EXCLUDED.has_format_conversion,
usage_routing_snapshots.has_format_conversion
),
candidate_id = CASE WHEN $14 THEN EXCLUDED.candidate_id ELSE COALESCE(EXCLUDED.candidate_id, usage_routing_snapshots.candidate_id) END,
candidate_index = CASE WHEN $14 THEN EXCLUDED.candidate_index ELSE COALESCE(EXCLUDED.candidate_index, usage_routing_snapshots.candidate_index) END,
key_name = CASE WHEN $14 THEN EXCLUDED.key_name ELSE COALESCE(EXCLUDED.key_name, usage_routing_snapshots.key_name) END,
planner_kind = CASE WHEN $14 THEN EXCLUDED.planner_kind ELSE COALESCE(EXCLUDED.planner_kind, usage_routing_snapshots.planner_kind) END,
route_family = CASE WHEN $14 THEN EXCLUDED.route_family ELSE COALESCE(EXCLUDED.route_family, usage_routing_snapshots.route_family) END,
route_kind = CASE WHEN $14 THEN EXCLUDED.route_kind ELSE COALESCE(EXCLUDED.route_kind, usage_routing_snapshots.route_kind) END,
execution_path = CASE WHEN $14 THEN EXCLUDED.execution_path ELSE COALESCE(EXCLUDED.execution_path, usage_routing_snapshots.execution_path) END,
local_execution_runtime_miss_reason = CASE WHEN $14 THEN EXCLUDED.local_execution_runtime_miss_reason ELSE COALESCE(EXCLUDED.local_execution_runtime_miss_reason, usage_routing_snapshots.local_execution_runtime_miss_reason) END,
selected_provider_id = CASE WHEN $14 THEN EXCLUDED.selected_provider_id ELSE COALESCE(EXCLUDED.selected_provider_id, usage_routing_snapshots.selected_provider_id) END,
selected_endpoint_id = CASE WHEN $14 THEN EXCLUDED.selected_endpoint_id ELSE COALESCE(EXCLUDED.selected_endpoint_id, usage_routing_snapshots.selected_endpoint_id) END,
selected_provider_api_key_id = CASE WHEN $14 THEN EXCLUDED.selected_provider_api_key_id ELSE COALESCE(EXCLUDED.selected_provider_api_key_id, usage_routing_snapshots.selected_provider_api_key_id) END,
has_format_conversion = CASE WHEN $14 THEN EXCLUDED.has_format_conversion ELSE COALESCE(EXCLUDED.has_format_conversion, usage_routing_snapshots.has_format_conversion) END,
updated_at = NOW()
crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_settlement_pricing_snapshot_sql.sql
+28 -108
View File
@@ -61,112 +61,32 @@ INSERT INTO usage_settlement_snapshots (
)
ON CONFLICT (request_id)
DO UPDATE SET
billing_snapshot_schema_version = COALESCE(
EXCLUDED.billing_snapshot_schema_version,
usage_settlement_snapshots.billing_snapshot_schema_version
),
billing_snapshot_status = COALESCE(
EXCLUDED.billing_snapshot_status,
usage_settlement_snapshots.billing_snapshot_status
),
settlement_snapshot_schema_version = COALESCE(
EXCLUDED.settlement_snapshot_schema_version,
usage_settlement_snapshots.settlement_snapshot_schema_version
),
settlement_snapshot = COALESCE(
EXCLUDED.settlement_snapshot,
usage_settlement_snapshots.settlement_snapshot
),
billing_dimensions = COALESCE(
EXCLUDED.billing_dimensions,
usage_settlement_snapshots.billing_dimensions
),
billing_input_tokens = COALESCE(
EXCLUDED.billing_input_tokens,
usage_settlement_snapshots.billing_input_tokens
),
billing_effective_input_tokens = COALESCE(
EXCLUDED.billing_effective_input_tokens,
usage_settlement_snapshots.billing_effective_input_tokens
),
billing_output_tokens = COALESCE(
EXCLUDED.billing_output_tokens,
usage_settlement_snapshots.billing_output_tokens
),
billing_cache_creation_tokens = COALESCE(
EXCLUDED.billing_cache_creation_tokens,
usage_settlement_snapshots.billing_cache_creation_tokens
),
billing_cache_creation_5m_tokens = COALESCE(
EXCLUDED.billing_cache_creation_5m_tokens,
usage_settlement_snapshots.billing_cache_creation_5m_tokens
),
billing_cache_creation_1h_tokens = COALESCE(
EXCLUDED.billing_cache_creation_1h_tokens,
usage_settlement_snapshots.billing_cache_creation_1h_tokens
),
billing_cache_read_tokens = COALESCE(
EXCLUDED.billing_cache_read_tokens,
usage_settlement_snapshots.billing_cache_read_tokens
),
billing_total_input_context = COALESCE(
EXCLUDED.billing_total_input_context,
usage_settlement_snapshots.billing_total_input_context
),
billing_cache_creation_cost_usd = COALESCE(
EXCLUDED.billing_cache_creation_cost_usd,
usage_settlement_snapshots.billing_cache_creation_cost_usd
),
billing_cache_read_cost_usd = COALESCE(
EXCLUDED.billing_cache_read_cost_usd,
usage_settlement_snapshots.billing_cache_read_cost_usd
),
billing_total_cost_usd = COALESCE(
EXCLUDED.billing_total_cost_usd,
usage_settlement_snapshots.billing_total_cost_usd
),
billing_actual_total_cost_usd = COALESCE(
EXCLUDED.billing_actual_total_cost_usd,
usage_settlement_snapshots.billing_actual_total_cost_usd
),
billing_pricing_source = COALESCE(
EXCLUDED.billing_pricing_source,
usage_settlement_snapshots.billing_pricing_source
),
billing_rule_id = COALESCE(
EXCLUDED.billing_rule_id,
usage_settlement_snapshots.billing_rule_id
),
billing_rule_version = COALESCE(
EXCLUDED.billing_rule_version,
usage_settlement_snapshots.billing_rule_version
),
rate_multiplier = COALESCE(
EXCLUDED.rate_multiplier,
usage_settlement_snapshots.rate_multiplier
),
is_free_tier = COALESCE(
EXCLUDED.is_free_tier,
usage_settlement_snapshots.is_free_tier
),
input_price_per_1m = COALESCE(
EXCLUDED.input_price_per_1m,
usage_settlement_snapshots.input_price_per_1m
),
output_price_per_1m = COALESCE(
EXCLUDED.output_price_per_1m,
usage_settlement_snapshots.output_price_per_1m
),
cache_creation_price_per_1m = COALESCE(
EXCLUDED.cache_creation_price_per_1m,
usage_settlement_snapshots.cache_creation_price_per_1m
),
cache_read_price_per_1m = COALESCE(
EXCLUDED.cache_read_price_per_1m,
usage_settlement_snapshots.cache_read_price_per_1m
),
price_per_request = COALESCE(
EXCLUDED.price_per_request,
usage_settlement_snapshots.price_per_request
),
billing_status = CASE WHEN $30 THEN EXCLUDED.billing_status ELSE usage_settlement_snapshots.billing_status END,
billing_snapshot_schema_version = CASE WHEN $30 THEN EXCLUDED.billing_snapshot_schema_version ELSE COALESCE(EXCLUDED.billing_snapshot_schema_version, usage_settlement_snapshots.billing_snapshot_schema_version) END,
billing_snapshot_status = CASE WHEN $30 THEN EXCLUDED.billing_snapshot_status ELSE COALESCE(EXCLUDED.billing_snapshot_status, usage_settlement_snapshots.billing_snapshot_status) END,
settlement_snapshot_schema_version = CASE WHEN $30 THEN EXCLUDED.settlement_snapshot_schema_version ELSE COALESCE(EXCLUDED.settlement_snapshot_schema_version, usage_settlement_snapshots.settlement_snapshot_schema_version) END,
settlement_snapshot = CASE WHEN $30 THEN EXCLUDED.settlement_snapshot ELSE COALESCE(EXCLUDED.settlement_snapshot, usage_settlement_snapshots.settlement_snapshot) END,
billing_dimensions = CASE WHEN $30 THEN EXCLUDED.billing_dimensions ELSE COALESCE(EXCLUDED.billing_dimensions, usage_settlement_snapshots.billing_dimensions) END,
billing_input_tokens = CASE WHEN $30 THEN EXCLUDED.billing_input_tokens ELSE COALESCE(EXCLUDED.billing_input_tokens, usage_settlement_snapshots.billing_input_tokens) END,
billing_effective_input_tokens = CASE WHEN $30 THEN EXCLUDED.billing_effective_input_tokens ELSE COALESCE(EXCLUDED.billing_effective_input_tokens, usage_settlement_snapshots.billing_effective_input_tokens) END,
billing_output_tokens = CASE WHEN $30 THEN EXCLUDED.billing_output_tokens ELSE COALESCE(EXCLUDED.billing_output_tokens, usage_settlement_snapshots.billing_output_tokens) END,
billing_cache_creation_tokens = CASE WHEN $30 THEN EXCLUDED.billing_cache_creation_tokens ELSE COALESCE(EXCLUDED.billing_cache_creation_tokens, usage_settlement_snapshots.billing_cache_creation_tokens) END,
billing_cache_creation_5m_tokens = CASE WHEN $30 THEN EXCLUDED.billing_cache_creation_5m_tokens ELSE COALESCE(EXCLUDED.billing_cache_creation_5m_tokens, usage_settlement_snapshots.billing_cache_creation_5m_tokens) END,
billing_cache_creation_1h_tokens = CASE WHEN $30 THEN EXCLUDED.billing_cache_creation_1h_tokens ELSE COALESCE(EXCLUDED.billing_cache_creation_1h_tokens, usage_settlement_snapshots.billing_cache_creation_1h_tokens) END,
billing_cache_read_tokens = CASE WHEN $30 THEN EXCLUDED.billing_cache_read_tokens ELSE COALESCE(EXCLUDED.billing_cache_read_tokens, usage_settlement_snapshots.billing_cache_read_tokens) END,
billing_total_input_context = CASE WHEN $30 THEN EXCLUDED.billing_total_input_context ELSE COALESCE(EXCLUDED.billing_total_input_context, usage_settlement_snapshots.billing_total_input_context) END,
billing_cache_creation_cost_usd = CASE WHEN $30 THEN EXCLUDED.billing_cache_creation_cost_usd ELSE COALESCE(EXCLUDED.billing_cache_creation_cost_usd, usage_settlement_snapshots.billing_cache_creation_cost_usd) END,
billing_cache_read_cost_usd = CASE WHEN $30 THEN EXCLUDED.billing_cache_read_cost_usd ELSE COALESCE(EXCLUDED.billing_cache_read_cost_usd, usage_settlement_snapshots.billing_cache_read_cost_usd) END,
billing_total_cost_usd = CASE WHEN $30 THEN EXCLUDED.billing_total_cost_usd ELSE COALESCE(EXCLUDED.billing_total_cost_usd, usage_settlement_snapshots.billing_total_cost_usd) END,
billing_actual_total_cost_usd = CASE WHEN $30 THEN EXCLUDED.billing_actual_total_cost_usd ELSE COALESCE(EXCLUDED.billing_actual_total_cost_usd, usage_settlement_snapshots.billing_actual_total_cost_usd) END,
billing_pricing_source = CASE WHEN $30 THEN EXCLUDED.billing_pricing_source ELSE COALESCE(EXCLUDED.billing_pricing_source, usage_settlement_snapshots.billing_pricing_source) END,
billing_rule_id = CASE WHEN $30 THEN EXCLUDED.billing_rule_id ELSE COALESCE(EXCLUDED.billing_rule_id, usage_settlement_snapshots.billing_rule_id) END,
billing_rule_version = CASE WHEN $30 THEN EXCLUDED.billing_rule_version ELSE COALESCE(EXCLUDED.billing_rule_version, usage_settlement_snapshots.billing_rule_version) END,
rate_multiplier = CASE WHEN $30 THEN EXCLUDED.rate_multiplier ELSE COALESCE(EXCLUDED.rate_multiplier, usage_settlement_snapshots.rate_multiplier) END,
is_free_tier = CASE WHEN $30 THEN EXCLUDED.is_free_tier ELSE COALESCE(EXCLUDED.is_free_tier, usage_settlement_snapshots.is_free_tier) END,
input_price_per_1m = CASE WHEN $30 THEN EXCLUDED.input_price_per_1m ELSE COALESCE(EXCLUDED.input_price_per_1m, usage_settlement_snapshots.input_price_per_1m) END,
output_price_per_1m = CASE WHEN $30 THEN EXCLUDED.output_price_per_1m ELSE COALESCE(EXCLUDED.output_price_per_1m, usage_settlement_snapshots.output_price_per_1m) END,
cache_creation_price_per_1m = CASE WHEN $30 THEN EXCLUDED.cache_creation_price_per_1m ELSE COALESCE(EXCLUDED.cache_creation_price_per_1m, usage_settlement_snapshots.cache_creation_price_per_1m) END,
cache_read_price_per_1m = CASE WHEN $30 THEN EXCLUDED.cache_read_price_per_1m ELSE COALESCE(EXCLUDED.cache_read_price_per_1m, usage_settlement_snapshots.cache_read_price_per_1m) END,
price_per_request = CASE WHEN $30 THEN EXCLUDED.price_per_request ELSE COALESCE(EXCLUDED.price_per_request, usage_settlement_snapshots.price_per_request) END,
updated_at = NOW()
crates/aether-data/adapters/postgres/src/usage/tests.rs
+655 -14
View File
@@ -1,25 +1,497 @@
use chrono::{TimeZone, Utc};
use serde_json::json;
use sqlx::Row;
use super::{
attach_compressed_body_refs, attach_usage_http_audit_body_refs,
attach_usage_routing_snapshot_metadata, attach_usage_settlement_pricing_snapshot_metadata,
inflate_usage_json_value, prepare_request_metadata_for_body_storage,
prepare_usage_body_storage, resolved_read_usage_body_ref, resolved_write_usage_body_ref,
clear_previous_request_body_facts, inflate_usage_json_value,
prepare_request_metadata_for_body_storage, prepare_usage_body_storage,
prepare_usage_upsert_context, request_body_capture_replaces_derived_facts,
resolved_read_usage_body_ref, resolved_write_usage_body_ref,
split_dashboard_daily_aggregate_range, split_dashboard_hourly_aggregate_range,
usage_body_capture_state_for_storage, usage_body_ref, usage_effective_input_tokens,
usage_http_audit_body_refs, usage_http_audit_capture_mode, usage_routing_snapshot_from_usage,
usage_settlement_pricing_snapshot_from_usage, usage_total_input_context, AggregateRangeSplit,
SqlxUsageReadRepository, UsageHttpAuditRefs, UsageRoutingSnapshot,
UsageSettlementPricingSnapshot, MAX_INLINE_USAGE_BODY_BYTES,
usage_body_capture_state_for_storage, usage_body_ref, usage_capture_update_allowed,
usage_effective_input_tokens, usage_http_audit_body_refs, usage_http_audit_capture_mode,
usage_routing_snapshot_from_usage, usage_settlement_pricing_snapshot_from_usage,
usage_total_input_context, AggregateRangeSplit, SqlxUsageReadRepository, UsageHttpAuditRefs,
UsageRoutingSnapshot, UsageSettlementPricingSnapshot, MAX_INLINE_USAGE_BODY_BYTES,
};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
use aether_data_contracts::repository::usage::{
UpsertUsageRecord, UsageBodyCaptureState, UsageBodyField, UsageCostSavingsSummaryQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageProviderPerformanceQuery,
UsageTimeSeriesGranularity,
UpsertUsageRecord, UsageAuditListQuery, UsageBodyCaptureState, UsageBodyField,
UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery,
UsageProviderPerformanceQuery, UsageTimeSeriesGranularity,
};
fn fast_clear_usage_record(
request_id: &str,
provider_name: &str,
now_unix_secs: u64,
terminal: bool,
terminal_state: UsageBodyCaptureState,
terminal_service_tier: Option<&str>,
) -> UpsertUsageRecord {
UpsertUsageRecord {
request_id: request_id.to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
provider_name: provider_name.to_string(),
model: "gpt-5".to_string(),
target_model: Some("gpt-5".to_string()),
provider_id: None,
provider_endpoint_id: None,
provider_api_key_id: None,
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
provider_api_family: Some("openai".to_string()),
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(false),
is_stream: Some(false),
input_tokens: terminal.then_some(1),
output_tokens: terminal.then_some(1),
total_tokens: terminal.then_some(2),
cache_creation_input_tokens: None,
cache_creation_ephemeral_5m_input_tokens: None,
cache_creation_ephemeral_1h_input_tokens: None,
cache_read_input_tokens: None,
cache_creation_cost_usd: None,
cache_read_cost_usd: None,
output_price_per_1m: None,
total_cost_usd: None,
actual_total_cost_usd: None,
status_code: terminal.then_some(200),
error_message: None,
error_category: None,
response_time_ms: terminal.then_some(10),
first_byte_time_ms: None,
status: if terminal { "completed" } else { "pending" }.to_string(),
billing_status: "pending".to_string(),
request_headers: None,
request_body: None,
request_body_ref: None,
request_body_state: None,
provider_request_headers: None,
provider_request_body: (!terminal).then(|| {
json!({
"model": "gpt-5",
"service_tier": "priority"
})
}),
provider_request_body_ref: None,
provider_request_body_state: Some(if terminal {
terminal_state
} else {
UsageBodyCaptureState::Inline
}),
response_headers: None,
response_body: None,
response_body_ref: None,
response_body_state: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
planner_kind: None,
route_family: None,
route_kind: None,
execution_path: None,
local_execution_runtime_miss_reason: None,
request_metadata: if terminal {
terminal_service_tier.map(|tier| json!({"provider_service_tier": tier}))
} else {
Some(json!({"provider_service_tier": "priority"}))
},
finalized_at_unix_secs: terminal.then_some(now_unix_secs + 1),
created_at_unix_ms: Some(now_unix_secs),
updated_at_unix_secs: now_unix_secs + u64::from(terminal),
}
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_terminal_none_capture_clears_fast_from_detail_and_lightweight_lists() {
let database_url = std::env::var("AETHER_TEST_DATABASE_URL")
.expect("AETHER_TEST_DATABASE_URL must point at the test database");
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url,
min_connections: 1,
max_connections: 2,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let repository =
SqlxUsageReadRepository::new(factory.connect_lazy().expect("lazy pool should build"));
crate::run_migrations(repository.pool())
.await
.expect("test database migrations should succeed");
let suffix = uuid::Uuid::new_v4().simple().to_string();
let request_id = format!("req-fast-clear-{suffix}");
let provider_name = format!("fast-clear-{suffix}");
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let mut pending_record = fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs,
false,
UsageBodyCaptureState::Inline,
Some("priority"),
);
pending_record.candidate_id = Some("candidate-a".to_string());
pending_record.candidate_index = Some(1);
pending_record.key_name = Some("key-a".to_string());
pending_record.planner_kind = Some("planner-a".to_string());
pending_record.route_family = Some("route-family-a".to_string());
pending_record.route_kind = Some("route-kind-a".to_string());
pending_record.execution_path = Some("path-a".to_string());
pending_record.local_execution_runtime_miss_reason = Some("miss-a".to_string());
pending_record.request_metadata = Some(json!({
"trace_id": "trace-a",
"provider_service_tier": "priority",
"provider_actual_service_tier": "priority",
"billing_snapshot": {
"schema_version": "2.0",
"status": "complete",
"resolved_variables": {
"input_price_per_1m": 30.0,
"output_price_per_1m": 150.0
}
},
"settlement_snapshot": {
"schema_version": "2.0",
"pricing_snapshot": {
"pricing_source": "processing_tier",
"service_tier": "priority"
},
"billing_plan_snapshot": {
"rule_id": "fast-rule",
"rule_version": "1"
}
},
"billing_dimensions": {"service_tier": "priority"},
"rate_multiplier": 2.0,
"input_price_per_1m": 30.0,
"output_price_per_1m": 150.0
}));
let pending = repository
.upsert(pending_record)
.await
.expect("pending usage should persist");
assert_eq!(pending.provider_service_tier().as_deref(), Some("priority"));
let mut terminal_record = fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs,
true,
UsageBodyCaptureState::None,
None,
);
terminal_record.provider_id = Some("final-provider-id".to_string());
terminal_record.provider_endpoint_id = Some("final-endpoint-id".to_string());
terminal_record.provider_api_key_id = Some("final-key-id".to_string());
terminal_record.target_model = None;
let terminal = repository
.upsert(terminal_record)
.await
.expect("terminal usage should persist");
assert_eq!(terminal.provider_service_tier(), None);
let stored = repository
.find_by_request_id(&request_id)
.await
.expect("detail lookup should succeed")
.expect("usage should exist");
assert_eq!(
stored.provider_request_body_state,
Some(UsageBodyCaptureState::None)
);
assert_eq!(stored.provider_service_tier(), None);
assert_eq!(stored.provider_actual_service_tier(), None);
let stored_metadata = stored
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("terminal audit metadata should remain");
assert_eq!(
stored_metadata
.get("trace_id")
.and_then(serde_json::Value::as_str),
Some("trace-a")
);
for stale_key in [
"provider_actual_service_tier",
"billing_snapshot",
"settlement_snapshot",
"billing_dimensions",
"rate_multiplier",
"input_price_per_1m",
"output_price_per_1m",
"candidate_id",
] {
assert!(
stored_metadata.get(stale_key).is_none(),
"terminal metadata retained stale key {stale_key}"
);
}
assert_eq!(stored.settlement_rate_multiplier(), None);
assert_eq!(stored.settlement_input_price_per_1m(), None);
assert_eq!(stored.settlement_output_price_per_1m(), None);
let settlement_row = sqlx::query(
"SELECT settlement_snapshot, billing_dimensions, CAST(rate_multiplier AS DOUBLE PRECISION) AS rate_multiplier, CAST(input_price_per_1m AS DOUBLE PRECISION) AS input_price_per_1m, CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m FROM usage_settlement_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("settlement snapshot row should be readable");
assert!(settlement_row
.try_get::<Option<serde_json::Value>, _>("settlement_snapshot")
.expect("settlement snapshot should decode")
.is_none());
assert!(settlement_row
.try_get::<Option<serde_json::Value>, _>("billing_dimensions")
.expect("billing dimensions should decode")
.is_none());
assert!(settlement_row
.try_get::<Option<f64>, _>("rate_multiplier")
.expect("rate multiplier should decode")
.is_none());
assert!(settlement_row
.try_get::<Option<f64>, _>("input_price_per_1m")
.expect("input price should decode")
.is_none());
assert!(settlement_row
.try_get::<Option<f64>, _>("output_price_per_1m")
.expect("output price should decode")
.is_none());
let physical_metadata = sqlx::query_scalar::<_, Option<serde_json::Value>>(
"SELECT request_metadata FROM \"usage\" WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("physical metadata should be readable")
.expect("clear tombstone should be stored");
assert!(physical_metadata.get("provider_service_tier").is_none());
let physical_body = sqlx::query(
"SELECT provider_request_body, provider_request_body_compressed FROM \"usage\" WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("physical body columns should be readable");
assert!(physical_body
.try_get::<Option<serde_json::Value>, _>("provider_request_body")
.expect("provider body column should decode")
.is_none());
assert!(physical_body
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
.expect("compressed provider body column should decode")
.is_none());
let physical_http = sqlx::query(
"SELECT provider_request_body_ref, provider_request_body_state FROM usage_http_audits WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("HTTP audit row should be readable");
assert!(physical_http
.try_get::<Option<String>, _>("provider_request_body_ref")
.expect("provider body ref should decode")
.is_none());
assert_eq!(
physical_http
.try_get::<Option<String>, _>("provider_request_body_state")
.expect("provider body state should decode")
.as_deref(),
Some("none")
);
let physical_blob_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = $1 AND body_field = 'provider_request_body'",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("provider body blob count should be readable");
assert_eq!(physical_blob_count, 0);
// A late pending candidate must not resurrect the terminal capture, metadata, or detached
// blob after the final state has been written.
let mut late_record = fast_clear_usage_record(
&request_id,
"late-provider",
now_unix_secs + 2,
false,
UsageBodyCaptureState::Inline,
Some("priority"),
);
late_record.model = "late-model".to_string();
late_record.target_model = Some("late-target".to_string());
late_record.provider_id = Some("late-provider-id".to_string());
late_record.provider_endpoint_id = Some("late-endpoint-id".to_string());
late_record.provider_api_key_id = Some("late-key-id".to_string());
late_record.endpoint_api_format = Some("late:format".to_string());
late_record.candidate_id = Some("late-candidate".to_string());
let late = repository
.upsert(late_record)
.await
.expect("late pending usage should be accepted without regressing capture");
assert_eq!(late.provider_service_tier(), None);
assert_eq!(late.provider_name, provider_name);
assert_eq!(late.model, "gpt-5");
assert_eq!(late.target_model, None);
assert_eq!(late.provider_id.as_deref(), Some("final-provider-id"));
assert_eq!(
late.provider_endpoint_id.as_deref(),
Some("final-endpoint-id")
);
assert_eq!(late.provider_api_key_id.as_deref(), Some("final-key-id"));
assert_eq!(late.endpoint_api_format.as_deref(), Some("openai:chat"));
assert_eq!(late.candidate_id, None);
let late_blob_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = $1 AND body_field = 'provider_request_body'",
)
.bind(&request_id)
.fetch_one(repository.pool())
.await
.expect("late provider body blob count should be readable");
assert_eq!(late_blob_count, 0);
let list_item = repository
.list_usage_audits(&UsageAuditListQuery {
provider_name: Some(provider_name.clone()),
limit: Some(10),
newest_first: true,
..UsageAuditListQuery::default()
})
.await
.expect("usage list should succeed")
.into_iter()
.find(|item| item.request_id == request_id)
.expect("usage list should contain the test row");
assert_eq!(list_item.provider_request_body_state, None);
assert_eq!(list_item.provider_service_tier(), None);
assert_eq!(list_item.target_model, None);
assert_eq!(list_item.candidate_id, None);
assert_eq!(list_item.candidate_index, None);
assert_eq!(list_item.key_name, None);
assert_eq!(list_item.planner_kind, None);
assert_eq!(list_item.route_family, None);
assert_eq!(list_item.route_kind, None);
assert_eq!(list_item.execution_path, None);
assert_eq!(list_item.local_execution_runtime_miss_reason, None);
let recent_item = repository
.list_recent_usage_audits(None, 20)
.await
.expect("recent usage list should succeed")
.into_iter()
.find(|item| item.request_id == request_id)
.expect("recent usage list should contain the test row");
assert_eq!(recent_item.provider_request_body_state, None);
assert_eq!(recent_item.provider_service_tier(), None);
sqlx::query("DELETE FROM \"usage\" WHERE request_id = $1")
.bind(&request_id)
.execute(repository.pool())
.await
.expect("test usage should be removed");
for (index, state, incoming_tier, expected_tier) in [
("disabled", UsageBodyCaptureState::Disabled, None, None),
("truncated", UsageBodyCaptureState::Truncated, None, None),
(
"unavailable",
UsageBodyCaptureState::Unavailable,
None,
None,
),
(
"none-stale-metadata",
UsageBodyCaptureState::None,
Some("priority"),
None,
),
(
"disabled-preserved",
UsageBodyCaptureState::Disabled,
Some("priority"),
Some("priority"),
),
(
"truncated-preserved",
UsageBodyCaptureState::Truncated,
Some("priority"),
Some("priority"),
),
] {
let request_id = format!("req-fast-clear-{suffix}-{index}");
let provider_name = format!("fast-clear-{suffix}-{index}");
repository
.upsert(fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs,
false,
UsageBodyCaptureState::Inline,
Some("priority"),
))
.await
.expect("typed-state pending usage should persist");
let terminal = repository
.upsert(fast_clear_usage_record(
&request_id,
&provider_name,
now_unix_secs + 1,
true,
state,
incoming_tier,
))
.await
.expect("typed-state terminal usage should persist");
assert_eq!(
terminal.provider_service_tier().as_deref(),
expected_tier,
"state={state:?} should use only same-source metadata"
);
let list_item = repository
.list_usage_audits(&UsageAuditListQuery {
provider_name: Some(provider_name.clone()),
limit: Some(10),
newest_first: true,
..UsageAuditListQuery::default()
})
.await
.expect("typed-state list should succeed")
.into_iter()
.find(|item| item.request_id == request_id)
.expect("typed-state row should be listed");
assert_eq!(list_item.provider_service_tier().as_deref(), expected_tier);
sqlx::query("DELETE FROM \"usage\" WHERE request_id = $1")
.bind(&request_id)
.execute(repository.pool())
.await
.expect("typed-state test usage should be removed");
}
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and a populated PostgreSQL database"]
async fn live_dashboard_combined_summary_matches_separate_queries() {
@@ -1536,6 +2008,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
assert!(sql.contains("request_metadata->>'client_ip'"));
assert!(sql.contains("'user_agent'"));
assert!(sql.contains("request_metadata->>'user_agent'"));
assert!(sql.contains("request_metadata->>'requested_reasoning_effort'"));
assert!(sql.contains("request_metadata->>'provider_reasoning_effort'"));
assert!(sql.contains("request_metadata->>'provider_service_tier'"));
assert!(sql.contains("request_metadata->>'provider_actual_service_tier'"));
assert!(sql.contains("AS client_family"));
@@ -1656,7 +2130,7 @@ fn usage_sql_settlement_pricing_snapshot_billing_values_use_authoritative_incomi
"billing_actual_total_cost_usd",
] {
let assignment = format!(
"{field} = COALESCE(\n EXCLUDED.{field},\n usage_settlement_snapshots.{field}\n )"
"{field} = CASE WHEN $30 THEN EXCLUDED.{field} ELSE COALESCE(EXCLUDED.{field}, usage_settlement_snapshots.{field}) END"
);
assert!(
sql.contains(assignment.as_str()),
@@ -1672,9 +2146,9 @@ fn usage_sql_settlement_pricing_snapshot_billing_values_use_authoritative_incomi
#[test]
fn usage_sql_upsert_recovers_missing_provider_links_after_billing_finalizes() {
for assignment in [
"provider_id = CASE WHEN \"usage\".billing_status = 'pending' OR (\"usage\".provider_id IS NULL AND (\"usage\".provider_endpoint_id IS NULL OR \"usage\".provider_endpoint_id = EXCLUDED.provider_endpoint_id) AND (\"usage\".provider_api_key_id IS NULL OR \"usage\".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_id, \"usage\".provider_id) ELSE \"usage\".provider_id END",
"provider_endpoint_id = CASE WHEN \"usage\".billing_status = 'pending' OR (\"usage\".provider_endpoint_id IS NULL AND (\"usage\".provider_id IS NULL OR \"usage\".provider_id = EXCLUDED.provider_id) AND (\"usage\".provider_api_key_id IS NULL OR \"usage\".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_endpoint_id, \"usage\".provider_endpoint_id) ELSE \"usage\".provider_endpoint_id END",
"provider_api_key_id = CASE WHEN \"usage\".billing_status = 'pending' OR (\"usage\".provider_api_key_id IS NULL AND (\"usage\".provider_id IS NULL OR \"usage\".provider_id = EXCLUDED.provider_id) AND (\"usage\".provider_endpoint_id IS NULL OR \"usage\".provider_endpoint_id = EXCLUDED.provider_endpoint_id)) THEN COALESCE(EXCLUDED.provider_api_key_id, \"usage\".provider_api_key_id) ELSE \"usage\".provider_api_key_id END",
"provider_id = CASE WHEN (\"usage\".billing_status = 'pending' AND $61) OR (\"usage\".billing_status <> 'pending' AND \"usage\".provider_id IS NULL AND (\"usage\".provider_endpoint_id IS NULL OR \"usage\".provider_endpoint_id = EXCLUDED.provider_endpoint_id) AND (\"usage\".provider_api_key_id IS NULL OR \"usage\".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_id, \"usage\".provider_id) ELSE \"usage\".provider_id END",
"provider_endpoint_id = CASE WHEN (\"usage\".billing_status = 'pending' AND $61) OR (\"usage\".billing_status <> 'pending' AND \"usage\".provider_endpoint_id IS NULL AND (\"usage\".provider_id IS NULL OR \"usage\".provider_id = EXCLUDED.provider_id) AND (\"usage\".provider_api_key_id IS NULL OR \"usage\".provider_api_key_id = EXCLUDED.provider_api_key_id)) THEN COALESCE(EXCLUDED.provider_endpoint_id, \"usage\".provider_endpoint_id) ELSE \"usage\".provider_endpoint_id END",
"provider_api_key_id = CASE WHEN (\"usage\".billing_status = 'pending' AND $61) OR (\"usage\".billing_status <> 'pending' AND \"usage\".provider_api_key_id IS NULL AND (\"usage\".provider_id IS NULL OR \"usage\".provider_id = EXCLUDED.provider_id) AND (\"usage\".provider_endpoint_id IS NULL OR \"usage\".provider_endpoint_id = EXCLUDED.provider_endpoint_id)) THEN COALESCE(EXCLUDED.provider_api_key_id, \"usage\".provider_api_key_id) ELSE \"usage\".provider_api_key_id END",
] {
assert!(
super::UPSERT_SQL.contains(assignment),
@@ -1761,6 +2235,56 @@ fn usage_sql_detached_body_flags_clear_inline_and_compressed_columns() {
.contains("WHEN EXCLUDED.client_response_body_compressed IS NOT NULL OR $60 THEN NULL"));
}
#[test]
fn usage_sql_capture_guard_covers_bodies_metadata_and_http_ref_tombstones() {
for assignment in [
"request_body = CASE WHEN \"usage\".billing_status = 'pending' AND $61",
"provider_request_body = CASE WHEN \"usage\".billing_status = 'pending' AND $61",
"response_body = CASE WHEN \"usage\".billing_status = 'pending' AND $61",
"client_response_body = CASE WHEN \"usage\".billing_status = 'pending' AND $61",
"request_metadata = CASE WHEN \"usage\".billing_status = 'pending' AND $61",
] {
assert!(
super::UPSERT_SQL.contains(assignment),
"missing guard: {assignment}"
);
}
for assignment in [
"WHEN EXCLUDED.request_body_state = 'none' THEN NULL",
"WHEN EXCLUDED.provider_request_body_state = 'none' THEN NULL",
"WHEN EXCLUDED.response_body_state = 'none' THEN NULL",
"WHEN EXCLUDED.client_response_body_state = 'none' THEN NULL",
] {
assert!(
super::UPSERT_USAGE_HTTP_AUDIT_SQL.contains(assignment),
"missing HTTP ref tombstone: {assignment}"
);
}
let source = include_str!("mod.rs");
assert!(source.contains(".bind(capture_update_allowed)"));
assert!(source.contains("if capture_update_allowed {"));
}
#[test]
fn usage_sql_terminal_snapshots_replace_sparse_routing_and_settlement_facts() {
assert!(super::UPSERT_SQL.contains(
"target_model = CASE WHEN \"usage\".billing_status = 'pending' AND $61 THEN CASE WHEN EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN EXCLUDED.target_model"
));
let routing = super::UPSERT_USAGE_ROUTING_SNAPSHOT_SQL;
assert!(routing.contains("candidate_id = CASE WHEN $14 THEN EXCLUDED.candidate_id"));
assert!(
routing.contains("selected_provider_id = CASE WHEN $14 THEN EXCLUDED.selected_provider_id")
);
let settlement = super::UPSERT_USAGE_SETTLEMENT_PRICING_SNAPSHOT_SQL;
assert!(settlement
.contains("settlement_snapshot = CASE WHEN $30 THEN EXCLUDED.settlement_snapshot"));
assert!(
settlement.contains("billing_dimensions = CASE WHEN $30 THEN EXCLUDED.billing_dimensions")
);
assert!(settlement.contains("rate_multiplier = CASE WHEN $30 THEN EXCLUDED.rate_multiplier"));
assert!(include_str!("mod.rs").contains(".bind(replace_existing)"));
}
#[test]
fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() {
assert!(super::UPSERT_SQL.contains(
@@ -1884,6 +2408,123 @@ fn usage_body_capture_state_for_storage_preserves_unavailable_states() {
);
}
#[test]
fn explicit_none_capture_replaces_stale_request_body_facts_even_with_a_residual_body() {
let stale_body = json!({
"reasoning_effort": "xhigh",
"service_tier": "priority"
});
assert!(request_body_capture_replaces_derived_facts(
Some(&stale_body),
Some(UsageBodyCaptureState::None),
));
assert!(request_body_capture_replaces_derived_facts(
Some(&stale_body),
None,
));
assert!(request_body_capture_replaces_derived_facts(
Some(&stale_body),
Some(UsageBodyCaptureState::Disabled),
));
assert!(request_body_capture_replaces_derived_facts(
None,
Some(UsageBodyCaptureState::Truncated),
));
assert!(!request_body_capture_replaces_derived_facts(None, None,));
}
#[test]
fn explicit_none_capture_drops_residual_body_ref_and_incoming_fast_metadata_before_storage() {
let mut usage = fast_clear_usage_record(
"req-none-residual",
"none-residual",
100,
true,
UsageBodyCaptureState::None,
Some("priority"),
);
usage.provider_request_body = Some(json!({
"model": "gpt-5",
"service_tier": "priority"
}));
usage.provider_request_body_ref =
Some("usage://request/req-none-residual/provider_request_body".to_string());
usage.request_metadata = Some(json!({
"trace_id": "trace-1",
"provider_service_tier": "priority",
"provider_reasoning_effort": "high",
"provider_cache_ttl_minutes": 30,
"provider_request_body_ref": "usage://request/req-none-residual/provider_request_body"
}));
let prepared = prepare_usage_upsert_context(&usage).expect("usage should prepare");
assert!(prepared.clear_provider_request_body);
assert!(!prepared.provider_request_body_storage.has_detached_blob());
assert_eq!(prepared.http_audit_refs.provider_request_body_ref, None);
assert_eq!(
prepared.request_metadata_value,
Some(json!({"trace_id": "trace-1"}))
);
}
#[test]
fn request_body_fact_clear_keeps_unrelated_metadata_and_emits_an_empty_tombstone() {
let previous = json!({
"trace_id": "trace-1",
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority",
"provider_cache_ttl_minutes": 30,
"provider_actual_service_tier": "default"
});
assert_eq!(
clear_previous_request_body_facts(Some(&previous), true, true),
json!({
"trace_id": "trace-1",
"provider_actual_service_tier": "default"
})
);
assert_eq!(
clear_previous_request_body_facts(
Some(&json!({"provider_service_tier": "priority"})),
false,
true,
),
json!({})
);
}
#[test]
fn terminal_capture_rejects_late_non_terminal_updates() {
assert!(usage_capture_update_allowed(None, "pending"));
assert!(usage_capture_update_allowed(
Some(("pending", "pending")),
"streaming",
));
assert!(usage_capture_update_allowed(
Some(("failed", "pending")),
"completed",
));
assert!(!usage_capture_update_allowed(
Some(("completed", "pending")),
"pending",
));
assert!(!usage_capture_update_allowed(
Some(("failed", "pending")),
"streaming",
));
assert!(!usage_capture_update_allowed(
Some(("streaming", "pending")),
"pending",
));
assert!(!usage_capture_update_allowed(
Some(("completed", "settled")),
"completed",
));
}
#[test]
fn prepare_request_metadata_for_body_storage_strips_body_ref_compatibility_keys() {
let detached = prepare_usage_body_storage(Some(&json!({