From 664c063a062726bb9e833e74cb29859f941fc43b Mon Sep 17 00:00:00 2001 From: elky Date: Fri, 17 Jul 2026 19:20:16 +0800 Subject: [PATCH] feat(usage): enrich audit metadata and detail views --- .../src/execution_runtime/stream/execution.rs | 159 ++++- .../stream/execution_failures.rs | 204 +++++- .../observability/monitoring/tests/trace.rs | 104 +++ .../admin/observability/monitoring/trace.rs | 26 +- .../admin/users/api_keys/responses/create.rs | 17 +- .../handlers/public/support/user_me_usage.rs | 27 + .../src/tests/control/admin/users.rs | 188 ++++- .../aether-admin/src/observability/usage.rs | 25 +- .../src/provider_compat/private_envelope.rs | 69 ++ crates/aether-billing/src/event_enrichment.rs | 160 +++-- crates/aether-billing/src/pricing.rs | 97 +-- crates/aether-billing/src/service.rs | 94 +-- .../aether-data/adapters/mysql/src/usage.rs | 50 +- .../adapters/mysql/src/usage/tests.rs | 44 ++ .../adapters/postgres/src/usage/mod.rs | 544 ++++++++++---- .../list_recent_usage_audits_prefix.sql | 3 + .../queries/list_usage_audits_prefix.sql | 3 + .../postgres/src/usage/queries/upsert_sql.sql | 50 +- .../queries/upsert_usage_http_audit_sql.sql | 49 +- .../upsert_usage_routing_snapshot_sql.sql | 42 +- ..._usage_settlement_pricing_snapshot_sql.sql | 136 +--- .../adapters/postgres/src/usage/tests.rs | 669 +++++++++++++++++- .../aether-data/adapters/sqlite/src/usage.rs | 50 +- .../adapters/sqlite/src/usage/tests.rs | 205 +++++- .../contracts/src/repository/usage/mod.rs | 9 +- .../contracts/src/repository/usage/types.rs | 257 ++++++- .../runtime/src/repository/usage/memory.rs | 323 +++++++-- .../src/repository/usage/memory/tests.rs | 123 +++- crates/aether-usage/runtime/src/record.rs | 194 ++++- .../runtime/src/request_metadata.rs | 279 +++++++- crates/aether-usage/runtime/src/runtime.rs | 172 ++++- crates/aether-usage/runtime/src/write.rs | 26 +- .../src/api/__tests__/provider-keys.spec.ts | 38 + .../api/__tests__/providers-summary.spec.ts | 101 +++ frontend/src/api/__tests__/users.spec.ts | 56 +- frontend/src/api/dashboard.ts | 2 + frontend/src/api/endpoints/keys.ts | 37 +- frontend/src/api/endpoints/providers.ts | 98 ++- frontend/src/api/me.ts | 2 + frontend/src/api/usage.ts | 2 + frontend/src/api/users.ts | 10 +- .../components/ProviderDetailDrawer.vue | 14 +- .../ModelMappingTab.contract.spec.ts | 38 + .../ProviderDetailDrawer.loading.spec.ts | 32 + .../provider-tabs/ModelMappingTab.vue | 17 +- .../usage/components/ElapsedTimeText.vue | 29 +- .../components/HorizontalRequestTimeline.vue | 42 +- .../usage/components/RequestDetailDrawer.vue | 345 +++++++-- .../usage/components/ServiceTierFacts.vue | 23 +- .../usage/components/UsageModelDisplay.vue | 169 +++++ .../usage/components/UsageRecordsTable.vue | 206 +----- .../__tests__/ElapsedTimeText.spec.ts | 40 +- .../HorizontalRequestTimeline.spec.ts | 89 ++- .../RequestDetailDrawer.pricing.spec.ts | 384 ++++++++++ .../__tests__/ServiceTierFacts.spec.ts | 25 +- .../__tests__/UsageRecordsTable.spec.ts | 161 ++++- .../__tests__/useUsageData.spec.ts | 166 ++++- .../usage/composables/useUsageData.ts | 130 +++- frontend/src/features/usage/types.ts | 3 +- .../usage/utils/__tests__/cyberError.spec.ts | 45 ++ .../usage/utils/__tests__/recordSync.spec.ts | 140 ++++ .../utils/__tests__/service-tier.spec.ts | 27 +- .../src/features/usage/utils/cyberError.ts | 70 ++ .../src/features/usage/utils/recordSync.ts | 158 ++++- .../src/features/usage/utils/service-tier.ts | 22 +- .../__tests__/apiKeyFeatureSettings.spec.ts | 69 ++ .../features/users/apiKeyFeatureSettings.ts | 64 ++ .../users/components/UserApiKeyFormDialog.vue | 37 +- .../provider-detail-contract.spec.ts | 72 ++ .../__tests__/usage-detail-contract.spec.ts | 52 ++ .../user-management-contract.spec.ts | 62 ++ frontend/src/mocks/handler.ts | 356 +++++++++- frontend/src/utils/featureSettings.ts | 9 + .../views/admin/AdminOperationsDashboard.vue | 59 +- frontend/src/views/admin/Users.vue | 119 +++- ...OperationsDashboard.manual-refresh.spec.ts | 19 +- .../__tests__/Users.request-loading.spec.ts | 46 ++ frontend/src/views/shared/Usage.vue | 159 +++-- .../Usage.initial-admin-load.spec.ts | 20 + 79 files changed, 6968 insertions(+), 1294 deletions(-) create mode 100644 frontend/src/api/__tests__/provider-keys.spec.ts create mode 100644 frontend/src/api/__tests__/providers-summary.spec.ts create mode 100644 frontend/src/features/providers/components/__tests__/ModelMappingTab.contract.spec.ts create mode 100644 frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts create mode 100644 frontend/src/features/usage/components/UsageModelDisplay.vue create mode 100644 frontend/src/features/usage/utils/__tests__/cyberError.spec.ts create mode 100644 frontend/src/features/usage/utils/cyberError.ts create mode 100644 frontend/src/features/users/__tests__/apiKeyFeatureSettings.spec.ts create mode 100644 frontend/src/features/users/apiKeyFeatureSettings.ts create mode 100644 frontend/src/mocks/__tests__/provider-detail-contract.spec.ts create mode 100644 frontend/src/mocks/__tests__/usage-detail-contract.spec.ts create mode 100644 frontend/src/mocks/__tests__/user-management-contract.spec.ts diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 220df410f..ecaa4f7a4 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -131,6 +131,7 @@ const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n"; const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024; const SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES: usize = 1024 * 1024; +const PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60); const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000; const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750); @@ -140,6 +141,78 @@ const DIRECT_PASSTHROUGH_CHANNEL_CAPACITY_ENV: &str = "AETHER_GATEWAY_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY"; const DIRECT_PASSTHROUGH_MODE_ENV: &str = "AETHER_GATEWAY_DIRECT_PASSTHROUGH_MODE"; +/// Retains the incomplete tail needed to recognize provider error events split across transport +/// chunks without retaining an unbounded copy of the stream. +#[derive(Default)] +struct ProviderStreamErrorInspection { + buffered: Vec, +} + +impl ProviderStreamErrorInspection { + fn observe(&mut self, report_context: Option<&Value>, chunk: &[u8]) -> Option { + if chunk.is_empty() { + return None; + } + if let Some(error_body) = extract_provider_private_stream_error_body(report_context, chunk) + { + return Some(error_body); + } + + self.append_rolling(chunk); + let error_body = extract_provider_private_stream_error_body(report_context, &self.buffered); + if error_body.is_none() { + self.trim_completed_sse_events(); + } + error_body + } + + fn append_rolling(&mut self, chunk: &[u8]) { + if chunk.len() >= PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES { + self.buffered.clear(); + self.buffered.extend_from_slice( + &chunk[chunk.len() - PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES..], + ); + return; + } + + let overflow = self + .buffered + .len() + .saturating_add(chunk.len()) + .saturating_sub(PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES); + if overflow > 0 { + self.buffered.drain(..overflow); + } + self.buffered.extend_from_slice(chunk); + } + + fn trim_completed_sse_events(&mut self) { + let Ok(text) = std::str::from_utf8(&self.buffered) else { + return; + }; + if !text.lines().any(|line| { + let line = line.trim_start(); + line.starts_with("event:") || line.starts_with("data:") || line.starts_with(':') + }) { + return; + } + + let lf_end = self + .buffered + .windows(2) + .rposition(|window| window == b"\n\n") + .map(|index| index + 2); + let crlf_end = self + .buffered + .windows(4) + .rposition(|window| window == b"\r\n\r\n") + .map(|index| index + 4); + if let Some(event_end) = lf_end.into_iter().chain(crlf_end).max() { + self.buffered.drain(..event_end); + } + } +} + struct StageElapsedGuard { stage: &'static str, started_at: Instant, @@ -1192,6 +1265,7 @@ struct DirectPassthroughFinalizerCore { stream_usage_report_context: Option, stream_usage_observer: Option, stream_usage_observer_buffered: Vec, + provider_error_inspection: ProviderStreamErrorInspection, provider_buffered_body: Vec, buffered_body: Vec, provider_body_truncated: bool, @@ -1318,10 +1392,10 @@ impl DirectPassthroughFinalizer { chunk.as_ref(), ); } - if let Some(error_body_json) = extract_provider_private_stream_error_body( - core.stream_usage_report_context.as_ref(), - chunk.as_ref(), - ) { + if let Some(error_body_json) = core + .provider_error_inspection + .observe(core.stream_usage_report_context.as_ref(), chunk.as_ref()) + { let error_status_code = resolve_local_sync_error_status_code(core.status_code, &error_body_json); core.terminal_failure = Some(build_stream_failure_from_provider_error_body( @@ -1486,6 +1560,7 @@ impl DirectPassthroughFinalizerCore { stream_usage_report_context, stream_usage_observer: _, stream_usage_observer_buffered: _, + provider_error_inspection: _, provider_buffered_body, buffered_body, provider_body_truncated, @@ -2207,6 +2282,7 @@ async fn execute_stream_from_direct_passthrough( stream_usage_report_context, stream_usage_observer, stream_usage_observer_buffered: Vec::new(), + provider_error_inspection: ProviderStreamErrorInspection::default(), provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, @@ -2286,6 +2362,7 @@ async fn execute_stream_from_direct_passthrough( .as_ref() .map(|_| StreamingStandardTerminalObserver::default()); let mut stream_usage_observer_buffered = Vec::new(); + let mut provider_error_inspection = ProviderStreamErrorInspection::default(); let mut provider_buffered_body = Vec::new(); let mut buffered_body = Vec::new(); let mut provider_body_truncated = false; @@ -2458,7 +2535,7 @@ async fn execute_stream_from_direct_passthrough( provider_chunk.as_ref(), ); } - let provider_private_error_body_json = extract_provider_private_stream_error_body( + let provider_private_error_body_json = provider_error_inspection.observe( stream_usage_report_context.as_ref(), provider_chunk.as_ref(), ); @@ -5007,6 +5084,7 @@ async fn execute_stream_from_frame_stream( .filter(|_| !sync_json_stream_bridge_active_for_report) .map(|_| StreamingStandardTerminalObserver::default()); let mut stream_usage_observer_buffered = Vec::new(); + let mut provider_error_inspection = ProviderStreamErrorInspection::default(); append_stream_capture_bytes( &mut provider_buffered_body, &provider_prefetched_body_for_report, @@ -5168,6 +5246,16 @@ async fn execute_stream_from_frame_stream( let replay_chunk = normalized_prefetched_chunk .as_deref() .unwrap_or(provider_prefetched_body_for_report.as_slice()); + if let Some(error_body_json) = provider_error_inspection + .observe(stream_usage_report_context.as_ref(), replay_chunk) + { + let error_status_code = + resolve_local_sync_error_status_code(status_code, &error_body_json); + terminal_failure = Some(build_stream_failure_from_provider_error_body( + error_status_code, + &error_body_json, + )); + } if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), @@ -5343,11 +5431,8 @@ async fn execute_stream_from_frame_stream( } else { chunk }; - let provider_private_error_body_json = - extract_provider_private_stream_error_body( - stream_usage_report_context.as_ref(), - &normalized_chunk, - ); + let provider_private_error_body_json = provider_error_inspection + .observe(stream_usage_report_context.as_ref(), &normalized_chunk); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), @@ -5508,11 +5593,8 @@ async fn execute_stream_from_frame_stream( { match normalizer.finish() { Ok(normalized_chunk) if !normalized_chunk.is_empty() => { - let provider_private_error_body_json = - extract_provider_private_stream_error_body( - stream_usage_report_context.as_ref(), - &normalized_chunk, - ); + let provider_private_error_body_json = provider_error_inspection + .observe(stream_usage_report_context.as_ref(), &normalized_chunk); if let (Some(observer), Some(report_context)) = ( stream_usage_observer.as_mut(), stream_usage_report_context.as_ref(), @@ -6114,7 +6196,7 @@ mod tests { stream_terminal_summary_missing_observed_finish, stream_terminal_summary_missing_observed_finish_with_requirement, stream_terminal_summary_represents_failure_with_requirement, - ClientVisibleStreamCompletionTracker, DirectPassthroughMode, + ClientVisibleStreamCompletionTracker, DirectPassthroughMode, ProviderStreamErrorInspection, }; use crate::control::GatewayControlDecision; use crate::stage_metrics::RequestStageTrace; @@ -6238,6 +6320,51 @@ mod tests { .observe_chunk(b"data: {\"type\":\"response.completed\",\"response\":{}}\r\n\r\n")); } + #[test] + fn provider_error_inspection_detects_response_failed_at_every_chunk_boundary() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"status\":\"in_progress\"}}\n\n", + "event: response.failed\n", + "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"type\":\"invalid_request\",\"message\":\"cyber policy rejected the request\",\"code\":\"cyber_policy_violation\",\"param\":\"input\"}}}\n\n", + ) + .as_bytes(); + + for split in 1..body.len() { + let mut inspection = ProviderStreamErrorInspection::default(); + let detected = inspection + .observe(None, &body[..split]) + .or_else(|| inspection.observe(None, &body[split..])) + .unwrap_or_else(|| panic!("response.failed was missed at byte split {split}")); + + assert_eq!( + detected.pointer("/error/code"), + Some(&json!("cyber_policy_violation")), + "string provider code changed at byte split {split}" + ); + assert_eq!( + detected.pointer("/error/param"), + Some(&json!("input")), + "provider error fields changed at byte split {split}" + ); + } + + let mut inspection = ProviderStreamErrorInspection::default(); + let mut detected = None; + for byte in body.chunks(1) { + if let Some(error_body) = inspection.observe(None, byte) { + detected = Some(error_body); + break; + } + } + let detected = detected.expect("byte-wise response.failed stream should be detected"); + assert_eq!( + detected.pointer("/error/code"), + Some(&json!("cyber_policy_violation")) + ); + assert_eq!(detected.pointer("/error/param"), Some(&json!("input"))); + } + fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot { aether_contracts::ProxySnapshot { enabled: Some(true), diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs index 0e9702420..97885eb86 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs @@ -20,10 +20,9 @@ use crate::execution_runtime::submission::{ use crate::log_ids::short_request_id; use crate::orchestration::{ apply_local_execution_effect, resolve_local_failover_analysis_for_attempt, - trace_upstream_response_body, with_upstream_response_report_context, - LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect, - LocalExecutionEffectContext, LocalHealthFailureEffect, LocalOAuthInvalidationEffect, - LocalPoolErrorEffect, + with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, + LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, + LocalOAuthInvalidationEffect, LocalPoolErrorEffect, }; use crate::request_candidate_runtime::record_report_request_candidate_status; use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context; @@ -36,6 +35,7 @@ pub(super) struct StreamFailureReport { pub(super) error_type: String, pub(super) error_message: String, extra_error_fields: Map, + provider_body_json: Option, } #[derive(Serialize)] @@ -54,20 +54,28 @@ struct StreamFailureBodyFields<'a> { } impl StreamFailureReport { - fn into_body_json(self) -> Value { + fn into_body_jsons(self) -> (Value, Option) { let Self { status_code, error_type, error_message, mut extra_error_fields, + provider_body_json, } = self; extra_error_fields.insert("type".to_string(), Value::String(error_type)); extra_error_fields.insert("message".to_string(), Value::String(error_message)); extra_error_fields.insert("code".to_string(), Value::from(status_code)); - Value::Object(Map::from_iter([( + let normalized_body = Value::Object(Map::from_iter([( "error".to_string(), Value::Object(extra_error_fields), - )])) + )])); + match provider_body_json { + Some(provider_body) if provider_body != normalized_body => { + (provider_body, Some(normalized_body)) + } + Some(provider_body) => (provider_body, None), + None => (normalized_body, None), + } } pub(super) fn to_json_string(&self) -> serde_json::Result { @@ -94,6 +102,7 @@ pub(super) fn build_stream_failure_report( error_type, error_message, extra_error_fields: Map::new(), + provider_body_json: None, } } @@ -124,6 +133,7 @@ pub(super) fn build_stream_failure_from_execution_error( error_type, error_message, extra_error_fields: error_object, + provider_body_json: None, } } @@ -150,6 +160,7 @@ pub(super) fn build_stream_failure_from_provider_error_body( error_type, error_message, extra_error_fields: Map::new(), + provider_body_json: Some(body_json.clone()), } } @@ -185,19 +196,33 @@ fn build_stream_failure_sync_payload( failure: StreamFailureReport, ) -> GatewaySyncReportRequest { let status_code = failure.status_code; - let body = trace_upstream_response_body(None, provider_buffered_body); + let (body, client_body) = failure.into_body_jsons(); + headers.retain(|name, _| { + !name.eq_ignore_ascii_case("content-encoding") + && !name.eq_ignore_ascii_case("content-length") + && !name.eq_ignore_ascii_case("content-type") + }); + headers.insert("content-type".to_string(), "application/json".to_string()); let report_context = with_upstream_response_report_context( report_context.as_ref(), status_code, Some(&headers), - body.as_ref(), + Some(&body), None, None, ) .or(report_context); - headers.remove("content-encoding"); - headers.remove("content-length"); - headers.insert("content-type".to_string(), "application/json".to_string()); + let report_context = report_context.map(|mut context| { + if let Some(object) = context.as_object_mut() { + let response_headers = serde_json::to_value(&headers).unwrap_or(Value::Null); + object.insert( + "provider_response_headers".to_string(), + response_headers.clone(), + ); + object.insert("client_response_headers".to_string(), response_headers); + } + context + }); GatewaySyncReportRequest { trace_id: trace_id.to_string(), @@ -205,8 +230,8 @@ fn build_stream_failure_sync_payload( report_context, status_code, headers, - body_json: Some(failure.into_body_json()), - client_body_json: None, + body_json: Some(body), + client_body_json: client_body, body_base64: (!provider_buffered_body.is_empty()) .then(|| base64::engine::general_purpose::STANDARD.encode(provider_buffered_body)), telemetry, @@ -218,8 +243,9 @@ fn stream_failure_body_field<'a>( field: &str, ) -> Option<&'a str> { payload - .body_json + .client_body_json .as_ref() + .or(payload.body_json.as_ref()) .and_then(|body_json| body_json.get("error")) .and_then(|value| value.get(field)) .and_then(Value::as_str) @@ -477,3 +503,151 @@ pub(super) async fn submit_midstream_stream_failure( ); } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use base64::Engine as _; + use serde_json::json; + + use super::{build_stream_failure_from_provider_error_body, build_stream_failure_sync_payload}; + + #[test] + fn midstream_failure_trace_uses_terminal_error_instead_of_buffered_sse() { + let provider_buffered_body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"instructions\":\"AGENTS.md secret prompt\",\"tools\":[{\"name\":\"update_plan\"}]}}\n\n", + "event: response.failed\n", + "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"type\":\"invalid_request\",\"message\":\"This content was flagged for possible cybersecurity risk.\",\"code\":\"cyber_policy_violation\",\"param\":\"input\",\"details\":{\"policy_category\":\"cybersecurity\",\"appeal_allowed\":true}}}}\n\n", + ) + .as_bytes(); + let terminal_error = aether_ai_formats::api::extract_provider_private_stream_error_body( + None, + provider_buffered_body, + ) + .expect("raw upstream SSE should expose its terminal provider error JSON"); + let failure = build_stream_failure_from_provider_error_body(400, &terminal_error); + + let payload = build_stream_failure_sync_payload( + "trace-cyber-policy", + "openai_responses_sync_error".to_string(), + Some(json!({"request_id": "request-cyber-policy"})), + BTreeMap::from([ + ("Content-Encoding".to_string(), "gzip".to_string()), + ("Content-Length".to_string(), "4096".to_string()), + ("Content-Type".to_string(), "text/event-stream".to_string()), + ( + "x-request-id".to_string(), + "req_usage-cyber-risk-demo".to_string(), + ), + ]), + None, + provider_buffered_body, + failure, + ); + + let trace_body = payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/upstream_response/body")) + .expect("candidate trace should include the terminal error body"); + assert_eq!( + trace_body, + payload.body_json.as_ref().expect("usage error body") + ); + assert_eq!(trace_body, &terminal_error); + assert_eq!(trace_body["error"]["type"], json!("invalid_request")); + assert_eq!( + trace_body["error"]["message"], + json!("This content was flagged for possible cybersecurity risk.") + ); + assert_eq!(trace_body["error"]["code"], json!("cyber_policy_violation")); + assert_eq!(trace_body["error"]["param"], json!("input")); + assert_eq!( + trace_body["error"]["details"], + json!({ + "policy_category": "cybersecurity", + "appeal_allowed": true + }) + ); + assert_eq!( + payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/upstream_response/headers/content-type")), + Some(&json!("application/json")) + ); + assert_eq!( + payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/upstream_response/headers/x-request-id")), + Some(&json!("req_usage-cyber-risk-demo")) + ); + assert_eq!( + payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/provider_response_headers/content-type")), + Some(&json!("application/json")) + ); + assert_eq!( + payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/client_response_headers/content-type")), + Some(&json!("application/json")) + ); + let trace_headers = payload + .report_context + .as_ref() + .and_then(|context| context.pointer("/upstream_response/headers")) + .and_then(serde_json::Value::as_object) + .expect("candidate trace should include terminal JSON headers"); + assert!(!trace_headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-encoding"))); + assert!(!trace_headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-length"))); + assert_eq!( + payload.headers.get("content-type").map(String::as_str), + Some("application/json") + ); + assert!(!payload + .headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-encoding"))); + assert!(!payload + .headers + .keys() + .any(|name| name.eq_ignore_ascii_case("content-length"))); + assert_eq!( + payload + .client_body_json + .as_ref() + .and_then(|body| body.pointer("/error/message")), + Some(&json!( + "This content was flagged for possible cybersecurity risk." + )) + ); + assert_eq!( + payload + .client_body_json + .as_ref() + .and_then(|body| body.pointer("/error/code")), + Some(&json!(400)) + ); + assert_ne!(payload.client_body_json.as_ref(), Some(&terminal_error)); + assert!(!trace_body.to_string().contains("AGENTS.md secret prompt")); + + let raw_capture = payload + .body_base64 + .as_deref() + .and_then(|body| base64::engine::general_purpose::STANDARD.decode(body).ok()) + .expect("raw provider stream should remain available for usage auditing"); + assert_eq!(raw_capture, provider_buffered_body); + assert!(String::from_utf8_lossy(&raw_capture).contains("AGENTS.md secret prompt")); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs b/apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs index 6261047e7..7b1d87fdf 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs @@ -733,6 +733,110 @@ async fn admin_monitoring_trace_request_exposes_failed_candidate_upstream_respon assert!(extra.get("provider_response").is_none()); } +#[tokio::test] +async fn admin_monitoring_trace_request_prefers_ref_backed_usage_response_body() { + let mut candidate = sample_candidate( + "cand-used", + "request-ref-body", + 0, + RequestCandidateStatus::Failed, + Some(101), + Some(33), + Some(400), + ); + candidate.extra_data = Some(json!({ + "upstream_response": { + "status_code": 400, + "headers": { + "content-type": "text/event-stream", + "x-request-id": "stale-request-like-body" + }, + "body": { + "model": "gpt-5.6-sol", + "input": [{"role": "user", "content": "request prompt"}] + } + } + })); + + let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate])); + let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider()], + vec![sample_endpoint()], + vec![sample_key()], + )); + let mut usage = sample_usage( + "request-ref-body", + "provider-1", + "OpenAI", + 0, + 0.0, + "failed", + Some(400), + 100, + ); + usage.candidate_id = Some("cand-used".to_string()); + usage.response_headers = Some(json!({ + "content-type": "application/json", + "x-request-id": "req_usage-cyber-risk-demo" + })); + usage.response_body = Some(json!({ + "error": { + "type": "invalid_request", + "message": "This content was flagged for possible cybersecurity risk.", + "code": 400 + } + })); + usage.response_body_state = Some(UsageBodyCaptureState::Reference); + let usage_repository = Arc::new(InMemoryUsageReadRepository::seed_with_detached_bodies( + vec![usage], + )); + let data_state = + crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( + request_candidates, + usage_repository, + ) + .with_provider_catalog_reader(provider_catalog); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests(data_state); + let context = request_context( + http::Method::GET, + "/api/admin/monitoring/trace/request-ref-body", + ); + + let response = local_monitoring_response(&state, &context) + .await + .expect("handler should not error") + .expect("route should be handled locally"); + + assert_eq!(response.status(), http::StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("body should read"); + let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse"); + let upstream_response = &payload["candidates"][0]["extra_data"]["upstream_response"]; + assert_eq!( + upstream_response["headers"], + json!({ + "content-type": "application/json", + "x-request-id": "req_usage-cyber-risk-demo" + }) + ); + assert_eq!( + upstream_response["body"]["error"], + json!({ + "type": "invalid_request", + "message": "This content was flagged for possible cybersecurity risk.", + "code": 400 + }) + ); + assert!(upstream_response["body"].get("input").is_none()); + assert_eq!( + upstream_response["body_ref"], + json!("usage://request/request-ref-body/response_body") + ); +} + #[tokio::test] async fn admin_monitoring_trace_request_decodes_connect_json_response_body_refs() { let mut candidate = sample_candidate( diff --git a/apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs b/apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs index 3eda959e0..1f5a39616 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs @@ -30,6 +30,25 @@ struct ResolvedAdminMonitoringTrace { usage: Option, } +async fn hydrate_admin_monitoring_trace_response_body( + state: &AdminAppState<'_>, + mut usage: StoredRequestUsageAudit, +) -> Result { + let is_error_node = !usage.status.eq_ignore_ascii_case("completed") + || usage + .status_code + .is_some_and(|status| !(200..300).contains(&status)); + let response_body_ref = if is_error_node && usage.response_body.is_none() { + usage.response_body_ref.clone() + } else { + None + }; + if let Some(body_ref) = response_body_ref.as_deref() { + usage.response_body = state.resolve_request_usage_body_ref(body_ref).await?; + } + Ok(usage) +} + pub(super) async fn build_admin_monitoring_trace_request_response( state: &AdminAppState<'_>, request_context: &AdminRequestContext<'_>, @@ -92,6 +111,10 @@ async fn resolve_admin_monitoring_trace( .read_request_usage_audit_shallow(request_id) .await .map_err(|err| GatewayError::Internal(err.to_string()))?; + let usage = match usage { + Some(usage) => Some(hydrate_admin_monitoring_trace_response_body(state, usage).await?), + None => None, + }; return Ok(Some(ResolvedAdminMonitoringTrace { trace, usage })); } @@ -102,9 +125,10 @@ async fn resolve_admin_monitoring_trace( .await .map_err(|err| GatewayError::Internal(err.to_string()))? { - usage_candidates.push(usage); + usage_candidates.push(hydrate_admin_monitoring_trace_response_body(state, usage).await?); } if let Some(usage) = state.find_request_usage_by_id(request_id).await? { + let usage = hydrate_admin_monitoring_trace_response_body(state, usage).await?; if !usage_candidates.iter().any(|item| item.id == usage.id) { usage_candidates.push(usage); } diff --git a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs index fe940dfda..dfe0cb37d 100644 --- a/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs +++ b/apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs @@ -35,13 +35,16 @@ pub(crate) async fn build_admin_create_user_api_key_response( let Some(user_id) = admin_user_id_from_api_keys_path(request_context.path()) else { return Ok(build_admin_users_bad_request_response("缺少 user_id")); }; - if state.find_user_auth_by_id(&user_id).await?.is_none() { + let Some(target_user) = state.find_user_auth_by_id(&user_id).await? else { return Ok(( http::StatusCode::NOT_FOUND, Json(json!({ "detail": "用户不存在" })), ) .into_response()); - } + }; + // The authenticated principal authorizes this admin operation, but ownership and inherited + // policies must always come from the user selected in the request path. + let target_user_id = target_user.id; let Some(request_body) = request_body else { return Ok(( @@ -148,7 +151,7 @@ pub(crate) async fn build_admin_create_user_api_key_response( let Some(created) = state .create_user_api_key(aether_data::repository::auth::CreateUserApiKeyRecord { - user_id: user_id.clone(), + user_id: target_user_id.clone(), api_key_id: uuid::Uuid::new_v4().to_string(), key_hash: hash_admin_user_api_key(&plaintext_key), key_encrypted: Some(key_encrypted), @@ -174,7 +177,11 @@ pub(crate) async fn build_admin_create_user_api_key_response( let created = if allowed_providers.is_some() { match state - .set_user_api_key_allowed_providers(&user_id, &created.api_key_id, allowed_providers) + .set_user_api_key_allowed_providers( + &target_user_id, + &created.api_key_id, + allowed_providers, + ) .await? { Some(updated) => updated, @@ -186,7 +193,7 @@ pub(crate) async fn build_admin_create_user_api_key_response( let created = if feature_settings.is_some() { match state .set_user_api_key_feature_settings( - &user_id, + &target_user_id, &created.api_key_id, feature_settings.clone(), ) diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index e66900885..ad3fe49ba 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -506,6 +506,9 @@ fn build_users_me_usage_record_payload( if let Some(reasoning_effort) = item.provider_reasoning_effort() { payload["reasoning_effort"] = json!(reasoning_effort); } + if let Some(requested_reasoning_effort) = item.requested_reasoning_effort() { + payload["requested_reasoning_effort"] = json!(requested_reasoning_effort); + } if let Some(service_tier) = item.provider_service_tier() { payload["service_tier"] = json!(service_tier); } @@ -577,6 +580,9 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_ if let Some(reasoning_effort) = item.provider_reasoning_effort() { payload["reasoning_effort"] = json!(reasoning_effort); } + if let Some(requested_reasoning_effort) = item.requested_reasoning_effort() { + payload["requested_reasoning_effort"] = json!(requested_reasoning_effort); + } if let Some(service_tier) = item.provider_service_tier() { payload["service_tier"] = json!(service_tier); } @@ -1655,6 +1661,27 @@ mod tests { assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6); } + #[test] + fn user_usage_payloads_expose_requested_and_provider_reasoning_mapping() { + let item = StoredRequestUsageAudit { + request_body: Some(json!({ + "reasoning": { "effort": "xhigh" } + })), + provider_request_body: Some(json!({ + "reasoning": { "effort": "max" } + })), + ..sample_usage("completed") + }; + + let record = build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false); + let active = build_users_me_usage_active_payload(&item); + + assert_eq!(record["requested_reasoning_effort"], "xhigh"); + assert_eq!(active["requested_reasoning_effort"], "xhigh"); + assert_eq!(record["reasoning_effort"], "max"); + assert_eq!(active["reasoning_effort"], "max"); + } + #[test] fn user_usage_active_override_uses_terminal_candidate_latency() { let candidate = sample_candidate( diff --git a/apps/aether-gateway/src/tests/control/admin/users.rs b/apps/aether-gateway/src/tests/control/admin/users.rs index 832b3d28d..75c8da4f2 100644 --- a/apps/aether-gateway/src/tests/control/admin/users.rs +++ b/apps/aether-gateway/src/tests/control/admin/users.rs @@ -20,7 +20,7 @@ use http::StatusCode; use serde_json::json; use super::super::{ - build_router_with_state, issue_test_admin_access_token, start_server, AppState, + build_router_with_state, hash_api_key, issue_test_admin_access_token, start_server, AppState, }; use crate::constants::{ GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER, @@ -1672,6 +1672,192 @@ async fn gateway_handles_admin_user_api_key_routes_locally_with_trusted_admin_pr upstream_handle.abort(); } +#[tokio::test] +async fn admin_created_user_key_inherits_target_user_policy_not_admin_policy() { + let mut admin_snapshot = sample_admin_api_key_snapshot("admin-user", "admin-seed-key"); + admin_snapshot.user_role = "admin".to_string(); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![ + (Some(hash_api_key("sk-admin-seed")), admin_snapshot), + ( + Some(hash_api_key("sk-target-seed")), + sample_admin_api_key_snapshot("target-user", "target-seed-key"), + ), + ])); + let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![ + sample_admin_user_with_role("admin-user", "admin", "admin@example.com", "admin"), + sample_admin_user_with_role("target-user", "user", "target@example.com", "target"), + ])); + + let admin_group = user_repository + .create_user_group(UpsertUserGroupRecord { + name: "Admin OpenAI".to_string(), + description: None, + priority: 10, + allowed_providers: Some(vec!["openai".to_string()]), + allowed_providers_mode: "specific".to_string(), + allowed_api_formats: Some(vec!["openai:responses".to_string()]), + allowed_api_formats_mode: "specific".to_string(), + allowed_models: Some(vec!["gpt-5.4".to_string()]), + allowed_models_mode: "specific".to_string(), + rate_limit: Some(100), + rate_limit_mode: "custom".to_string(), + }) + .await + .expect("admin group should create") + .expect("admin group should exist"); + user_repository + .add_user_to_group(&admin_group.id, "admin-user") + .await + .expect("admin group membership should create"); + + let target_group = user_repository + .create_user_group(UpsertUserGroupRecord { + name: "Target Claude".to_string(), + description: None, + priority: 10, + allowed_providers: Some(vec!["anthropic".to_string()]), + allowed_providers_mode: "specific".to_string(), + allowed_api_formats: Some(vec!["claude:messages".to_string()]), + allowed_api_formats_mode: "specific".to_string(), + allowed_models: Some(vec!["claude-sonnet-4-5".to_string()]), + allowed_models_mode: "specific".to_string(), + rate_limit: Some(30), + rate_limit_mode: "custom".to_string(), + }) + .await + .expect("target group should create") + .expect("target group should exist"); + user_repository + .add_user_to_group(&target_group.id, "target-user") + .await + .expect("target group membership should create"); + user_repository + .update_user_feature_settings( + "admin-user", + Some(json!({"chat_pii_redaction": {"enabled": false}})), + ) + .await + .expect("admin feature settings should update"); + user_repository + .update_user_feature_settings( + "target-user", + Some(json!({"chat_pii_redaction": {"enabled": true}})), + ) + .await + .expect("target feature settings should update"); + + let data_state = GatewayDataState::with_auth_api_key_repository_for_tests(auth_repository) + .with_user_reader(user_repository); + let app_state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(data_state); + let inspection_state = app_state.clone(); + let gateway = build_router_with_state(app_state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!( + "{gateway_url}/api/admin/users/target-user/api-keys" + )) + .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "admin-session") + .json(&json!({"name": "target-key"})) + .send() + .await + .expect("request should succeed"); + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("response should parse"); + let plaintext_key = payload["key"].as_str().expect("plaintext key should exist"); + let api_key_id = payload["id"].as_str().expect("key id should exist"); + + let resolved = inspection_state + .read_cached_auth_api_key_snapshot_by_key_hash( + &hash_api_key(plaintext_key), + chrono::Utc::now().timestamp().max(0) as u64, + ) + .await + .expect("created key snapshot should resolve") + .expect("created key snapshot should exist"); + assert_eq!(resolved.user_id, "target-user"); + assert_eq!( + resolved.effective_allowed_providers(), + Some(&["anthropic".to_string()][..]) + ); + assert_eq!( + resolved.effective_allowed_api_formats(), + Some(&["claude:messages".to_string()][..]) + ); + assert_eq!( + resolved.effective_allowed_models(), + Some(&["claude-sonnet-4-5".to_string()][..]) + ); + assert_eq!(resolved.user_rate_limit, Some(30)); + + inspection_state + .update_user_group( + &target_group.id, + UpsertUserGroupRecord { + name: "Target Gemini".to_string(), + description: None, + priority: 10, + allowed_providers: Some(vec!["google".to_string()]), + allowed_providers_mode: "specific".to_string(), + allowed_api_formats: Some(vec!["gemini:generate-content".to_string()]), + allowed_api_formats_mode: "specific".to_string(), + allowed_models: Some(vec!["gemini-2.5-pro".to_string()]), + allowed_models_mode: "specific".to_string(), + rate_limit: Some(15), + rate_limit_mode: "custom".to_string(), + }, + ) + .await + .expect("target group should update") + .expect("target group should still exist"); + let updated = inspection_state + .read_cached_auth_api_key_snapshot_by_key_hash( + &hash_api_key(plaintext_key), + chrono::Utc::now().timestamp().max(0) as u64, + ) + .await + .expect("created key snapshot should refresh") + .expect("created key snapshot should still exist"); + assert_eq!(updated.user_id, "target-user"); + assert_eq!( + updated.effective_allowed_providers(), + Some(&["google".to_string()][..]) + ); + assert_eq!( + updated.effective_allowed_api_formats(), + Some(&["gemini:generate-content".to_string()][..]) + ); + assert_eq!( + updated.effective_allowed_models(), + Some(&["gemini-2.5-pro".to_string()][..]) + ); + assert_eq!(updated.user_rate_limit, Some(15)); + + let target_features = inspection_state + .read_user_feature_settings("target-user") + .await + .expect("target feature settings should resolve"); + let key_features = inspection_state + .read_auth_api_key_feature_settings("target-user", api_key_id, false) + .await + .expect("key feature settings should resolve"); + assert_eq!( + target_features, + Some(json!({"chat_pii_redaction": {"enabled": true}})) + ); + assert_eq!( + key_features, None, + "an omitted key override must inherit target settings" + ); + + gateway_handle.abort(); +} + #[tokio::test] async fn gateway_returns_conflict_for_admin_create_user_api_key_when_writer_unavailable() { let upstream_hits = Arc::new(Mutex::new(0usize)); diff --git a/crates/aether-admin/src/observability/usage.rs b/crates/aether-admin/src/observability/usage.rs index e19e81f12..e703fdcc0 100644 --- a/crates/aether-admin/src/observability/usage.rs +++ b/crates/aether-admin/src/observability/usage.rs @@ -1231,6 +1231,9 @@ fn admin_usage_active_request_json( if let Some(reasoning_effort) = item.provider_reasoning_effort() { value["reasoning_effort"] = json!(reasoning_effort); } + if let Some(requested_reasoning_effort) = item.requested_reasoning_effort() { + value["requested_reasoning_effort"] = json!(requested_reasoning_effort); + } if let Some(service_tier) = item.provider_service_tier() { value["service_tier"] = json!(service_tier); } @@ -1298,6 +1301,7 @@ pub fn admin_usage_record_json( "response_time_ms": item.response_time_ms, "first_byte_time_ms": item.first_byte_time_ms, "created_at": unix_secs_to_rfc3339(item.created_at_unix_ms), + "updated_at": unix_secs_to_rfc3339(item.updated_at_unix_secs), "input_price_per_1m": input_price_per_1m, "output_price_per_1m": output_price_per_1m, "cache_creation_price_per_1m": cache_creation_price_per_1m, @@ -1354,6 +1358,12 @@ pub fn admin_usage_record_json( if let Some(reasoning_effort) = item.provider_reasoning_effort() { object.insert("reasoning_effort".to_string(), json!(reasoning_effort)); } + if let Some(requested_reasoning_effort) = item.requested_reasoning_effort() { + object.insert( + "requested_reasoning_effort".to_string(), + json!(requested_reasoning_effort), + ); + } if let Some(service_tier) = item.provider_service_tier() { object.insert("service_tier".to_string(), json!(service_tier)); } @@ -2702,10 +2712,13 @@ mod tests { } #[test] - fn admin_usage_record_includes_provider_reasoning_effort() { + fn admin_usage_record_includes_requested_and_provider_reasoning_efforts() { let item = StoredRequestUsageAudit { + request_body: Some(json!({ + "reasoning": { "effort": "xhigh" } + })), provider_request_body: Some(json!({ - "reasoning": { "effort": "xhigh" }, + "reasoning": { "effort": "max" }, "service_tier": "priority" })), ..sample_usage("completed", Some(200), None) @@ -2721,10 +2734,14 @@ mod tests { ); let active = admin_usage_active_request_json(&item, None, None, None); - assert_eq!(record["reasoning_effort"], "xhigh"); - assert_eq!(active["reasoning_effort"], "xhigh"); + assert_eq!(record["requested_reasoning_effort"], "xhigh"); + assert_eq!(active["requested_reasoning_effort"], "xhigh"); + assert_eq!(record["reasoning_effort"], "max"); + assert_eq!(active["reasoning_effort"], "max"); assert_eq!(record["service_tier"], "priority"); assert_eq!(active["service_tier"], "priority"); + assert!(record["updated_at"].is_string()); + assert_eq!(record["updated_at"], active["updated_at"]); } #[test] diff --git a/crates/aether-ai/formats/src/provider_compat/private_envelope.rs b/crates/aether-ai/formats/src/provider_compat/private_envelope.rs index 1b5365d64..4eb01ad0f 100644 --- a/crates/aether-ai/formats/src/provider_compat/private_envelope.rs +++ b/crates/aether-ai/formats/src/provider_compat/private_envelope.rs @@ -2,6 +2,7 @@ use std::collections::BTreeMap; use serde_json::Value; +use crate::formats::shared::stream_core::common::openai_stream_terminal_error_body; use crate::formats::shared::AiSurfaceFinalizeError; use crate::provider_compat::kiro_stream::KiroToClaudeCliStreamState; @@ -719,6 +720,27 @@ fn extract_stream_error_event_body(body: &[u8]) -> Option { { return Some(normalize_provider_private_error_body(event)); } + if let Some(mut error_body) = openai_stream_terminal_error_body(&event) { + let response_failed_without_provider_type = event + .get("type") + .and_then(Value::as_str) + .is_some_and(|event_type| event_type.eq_ignore_ascii_case("response.failed")) + && event + .pointer("/response/error/type") + .or_else(|| event.pointer("/error/type")) + .and_then(Value::as_str) + .map(str::trim) + .is_none_or(str::is_empty); + if response_failed_without_provider_type { + if let Some(error) = error_body.get_mut("error").and_then(Value::as_object_mut) { + error.insert( + "type".to_string(), + Value::String("stream_terminal_error".to_string()), + ); + } + } + return Some(error_body); + } current_event_type = None; } None @@ -1179,6 +1201,53 @@ data: {"message":"bad"} assert!(stream_body_contains_error_event(body)); } + #[test] + fn extracts_openai_response_failed_error_without_losing_provider_fields() { + let body = br#"event: response.failed +data: {"type":"response.failed","response":{"status":"failed","error":{"type":"invalid_request","message":"cyber policy rejected the request","code":"cyber_policy_violation","param":"input","details":{"policy_category":"cybersecurity","appeal_allowed":true}}}} + +"#; + + let error_body = extract_provider_private_stream_error_body(None, body) + .expect("response.failed should expose its provider error body"); + + assert_eq!( + error_body, + json!({ + "error": { + "type": "invalid_request", + "message": "cyber policy rejected the request", + "code": "cyber_policy_violation", + "param": "input", + "details": { + "policy_category": "cybersecurity", + "appeal_allowed": true + } + } + }) + ); + } + + #[test] + fn response_failed_without_provider_type_uses_stream_terminal_error_category() { + let body = br#"event: response.failed +data: {"type":"response.failed","response":{"status":"failed","error":{"message":"cyber policy rejected the request","code":"cyber_policy"}}} + +"#; + + let error_body = extract_provider_private_stream_error_body(None, body) + .expect("response.failed should expose its provider error body"); + + assert_eq!( + error_body.pointer("/error/type"), + Some(&json!("stream_terminal_error")) + ); + assert_eq!( + error_body.pointer("/error/code"), + Some(&json!("cyber_policy")) + ); + } + #[test] fn windsurf_sync_error_message_is_not_normalized_as_success() { let report_context = json!({ diff --git a/crates/aether-billing/src/event_enrichment.rs b/crates/aether-billing/src/event_enrichment.rs index d25668854..0c98003a1 100644 --- a/crates/aether-billing/src/event_enrichment.rs +++ b/crates/aether-billing/src/event_enrichment.rs @@ -1,9 +1,7 @@ use aether_data_contracts::repository::billing::StoredBillingModelContext; use aether_data_contracts::repository::usage::{ - extract_provider_actual_service_tier_from_response, - extract_provider_cache_ttl_minutes_from_metadata, extract_provider_service_tier_from_body, - normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, - PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, + extract_provider_cache_ttl_minutes_from_metadata, resolve_provider_cache_ttl_minutes, + resolve_provider_service_tier_from_request_capture, }; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{UsageEvent, UsageEventType}; @@ -160,7 +158,9 @@ fn calculate_billing_computation( .clone() .or_else(|| event.data.api_format.clone()), requested_processing_tier: processing_tiers.requested, - actual_processing_tier: processing_tiers.actual, + // The response-reported tier remains usage audit data, but it is not authoritative for + // pricing. Settlement follows the final request that was sent upstream. + actual_processing_tier: None, request_count, input_tokens: event.data.input_tokens.unwrap_or_default() as i64, output_tokens: event.data.output_tokens.unwrap_or_default() as i64, @@ -192,29 +192,18 @@ fn calculate_billing_computation( #[derive(Debug, Clone, PartialEq, Eq)] struct UsageEventProcessingTiers { requested: Option, - actual: Option, } fn usage_event_processing_tiers( data: &aether_usage_runtime::UsageEventData, ) -> UsageEventProcessingTiers { - let metadata = data.request_metadata.as_ref().and_then(Value::as_object); - let requested = extract_provider_service_tier_from_body(data.provider_request_body.as_ref()) - .or_else(|| { - metadata - .and_then(|metadata| metadata.get(PROVIDER_SERVICE_TIER_METADATA_KEY)) - .and_then(Value::as_str) - .and_then(normalize_provider_service_tier) - }); - let actual = metadata - .and_then(|metadata| metadata.get(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY)) - .and_then(Value::as_str) - .and_then(normalize_provider_service_tier) - .or_else(|| { - extract_provider_actual_service_tier_from_response(data.response_body.as_ref()) - }); + let requested = resolve_provider_service_tier_from_request_capture( + data.provider_request_body.as_ref(), + data.provider_request_body_state, + data.request_metadata.as_ref(), + ); - UsageEventProcessingTiers { requested, actual } + UsageEventProcessingTiers { requested } } fn usage_event_provider_cache_ttl_minutes( @@ -406,6 +395,7 @@ fn build_settlement_snapshot( #[cfg(test)] mod tests { use aether_data_contracts::repository::billing::StoredBillingModelContext; + use aether_data_contracts::repository::usage::UsageBodyCaptureState; use aether_usage_runtime::{UsageEvent, UsageEventData, UsageEventType}; use async_trait::async_trait; use serde_json::json; @@ -444,7 +434,7 @@ mod tests { } #[test] - fn processing_tier_facts_keep_request_and_terminal_response_independent() { + fn processing_tier_fact_uses_final_provider_request_body() { let data = UsageEventData { provider_request_body: Some(json!({"service_tier": "Priority"})), response_body: Some(json!({"service_tier": "priority"})), @@ -458,7 +448,6 @@ mod tests { let tiers = usage_event_processing_tiers(&data); assert_eq!(tiers.requested.as_deref(), Some("priority")); - assert_eq!(tiers.actual.as_deref(), Some("default")); } #[test] @@ -474,7 +463,83 @@ mod tests { let tiers = usage_event_processing_tiers(&data); assert_eq!(tiers.requested.as_deref(), Some("fast")); - assert_eq!(tiers.actual.as_deref(), Some("fast")); + } + + #[test] + fn processing_tier_does_not_fall_back_to_stale_metadata_when_body_is_present() { + let data = UsageEventData { + provider_request_body: Some(json!({"model": "gpt-5"})), + response_body: Some(json!({"service_tier": "priority"})), + request_metadata: Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "priority" + })), + ..UsageEventData::default() + }; + + let tiers = usage_event_processing_tiers(&data); + + assert_eq!(tiers.requested, None); + } + + #[test] + fn processing_tier_uses_request_derived_metadata_when_body_capture_was_disabled() { + let data = UsageEventData { + provider_request_body: None, + provider_request_body_state: Some(UsageBodyCaptureState::Disabled), + response_body: Some(json!({"service_tier": "flex"})), + request_metadata: Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "flex" + })), + ..UsageEventData::default() + }; + + let tiers = usage_event_processing_tiers(&data); + + assert_eq!(tiers.requested.as_deref(), Some("priority")); + } + + #[test] + fn processing_tier_does_not_use_metadata_or_response_when_final_request_body_is_missing() { + let data = UsageEventData { + provider_request_body: None, + provider_request_body_state: Some(UsageBodyCaptureState::None), + response_body: Some(json!({"service_tier": "priority"})), + request_metadata: Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "priority" + })), + ..UsageEventData::default() + }; + + let tiers = usage_event_processing_tiers(&data); + + assert_eq!(tiers.requested, None); + } + + #[test] + fn processing_tier_uses_request_derived_metadata_when_body_capture_was_truncated() { + let data = UsageEventData { + provider_request_body: Some(json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded", + "max_bytes": 128, + "source_bytes": 4096, + "value_kind": "object" + })), + provider_request_body_state: Some(UsageBodyCaptureState::Truncated), + response_body: Some(json!({"service_tier": "default"})), + request_metadata: Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "default" + })), + ..UsageEventData::default() + }; + + let tiers = usage_event_processing_tiers(&data); + + assert_eq!(tiers.requested.as_deref(), Some("priority")); } #[tokio::test] @@ -638,7 +703,7 @@ mod tests { } #[tokio::test] - async fn settlement_uses_actual_processing_tier_catalog_and_source() { + async fn settlement_uses_requested_processing_tier_catalog_and_ignores_response_tier() { let lookup = TestLookup { name_context: Some( StoredBillingModelContext::new( @@ -706,13 +771,28 @@ mod tests { .and_then(|value| value.pointer("/settlement_snapshot/pricing_snapshot")) .expect("settlement pricing snapshot should exist"); assert_eq!(pricing_snapshot["requested_processing_tier"], "priority"); - assert_eq!(pricing_snapshot["actual_processing_tier"], "flex"); - assert_eq!(pricing_snapshot["billing_processing_tier"], "flex"); - assert_eq!(pricing_snapshot["tiered_pricing_source"], "global_default"); - assert_eq!(pricing_snapshot["processing_tier_price_multiplier"], 0.5); + assert!(pricing_snapshot["actual_processing_tier"].is_null()); + assert_eq!(pricing_snapshot["billing_processing_tier"], "priority"); + assert_eq!( + pricing_snapshot["tiered_pricing_source"], + "provider_override" + ); + assert_eq!( + pricing_snapshot["processing_tier_price_multiplier"], + Value::Null + ); assert_eq!( pricing_snapshot["tiered_pricing"]["tiers"][0]["input_price_per_1m"], - 2.5 + 10.0 + ); + // The response fact remains available for audit, but does not influence settlement. + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|value| value.get("provider_actual_service_tier")), + Some(&json!("flex")) ); assert_eq!( event @@ -720,12 +800,12 @@ mod tests { .request_metadata .as_ref() .and_then(|value| { value.pointer("/billing_dimensions/actual_processing_tier") }), - Some(&json!("flex")) + Some(&Value::Null) ); } #[tokio::test] - async fn actual_processing_catalog_controls_image_price_with_independent_fixed_price() { + async fn requested_processing_catalog_controls_image_price_with_independent_fixed_price() { let lookup = TestLookup { name_context: Some( StoredBillingModelContext::new( @@ -783,27 +863,27 @@ mod tests { .await .expect("billing should succeed"); - assert_eq!(event.data.total_cost_usd, Some(0.44)); - assert_eq!(event.data.actual_total_cost_usd, Some(0.44)); + assert_eq!(event.data.total_cost_usd, Some(0.84)); + assert_eq!(event.data.actual_total_cost_usd, Some(0.84)); let metadata = event.data.request_metadata.as_ref().expect("metadata"); let pricing = metadata .pointer("/settlement_snapshot/pricing_snapshot") .expect("pricing snapshot"); - assert_eq!(pricing["billing_processing_tier"], "flex"); - assert_eq!(pricing["tiered_pricing_source"], "global_default"); + assert_eq!(pricing["billing_processing_tier"], "priority"); + assert_eq!(pricing["tiered_pricing_source"], "provider_override"); assert_eq!(pricing["price_per_request_source"], "provider_override"); - assert_eq!(pricing["pricing_source"], "mixed"); + assert_eq!(pricing["pricing_source"], "provider_override"); assert_eq!( metadata .pointer("/billing_snapshot/resolved_variables/image_output_price_per_image") .and_then(Value::as_f64), - Some(0.2) + Some(0.4) ); assert_eq!( metadata .pointer("/billing_snapshot/cost_breakdown/image_output_cost") .and_then(Value::as_f64), - Some(0.4) + Some(0.8) ); assert_eq!( metadata diff --git a/crates/aether-billing/src/pricing.rs b/crates/aether-billing/src/pricing.rs index e00ce297b..d5cf0ff10 100644 --- a/crates/aether-billing/src/pricing.rs +++ b/crates/aether-billing/src/pricing.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeSet; - use aether_data_contracts::repository::{ billing::StoredBillingModelContext, global_models::{explicit_pricing_catalog_state, ExplicitPricingCatalogState}, @@ -108,17 +106,12 @@ impl BillingModelPricingSnapshot { ) -> BillingPricingResolution { let requested_processing_tier = normalize_processing_tier(requested_processing_tier); let actual_processing_tier = normalize_processing_tier(actual_processing_tier); - let billing_processing_tier = actual_processing_tier + // The provider-reported actual tier is retained in the resolution for audit only. The + // final upstream request is the sole authority for selecting a billing catalog. + let billing_processing_tier = requested_processing_tier .as_deref() .map(canonical_processing_tier) - .or_else(|| { - requested_processing_tier.as_deref().map_or( - Some("standard".to_string()), - |requested| { - processing_tier_is_standard(requested).then(|| "standard".to_string()) - }, - ) - }); + .or_else(|| Some("standard".to_string())); let (tiered_pricing, tiered_pricing_source, processing_tier_price_multiplier) = billing_processing_tier @@ -151,17 +144,12 @@ impl BillingModelPricingSnapshot { self.validate_processing_tier_containers()?; let requested_processing_tier = normalize_processing_tier(requested_processing_tier); let actual_processing_tier = normalize_processing_tier(actual_processing_tier); - let billing_processing_tier = actual_processing_tier + // Keep this checked path aligned with `resolve_pricing`: response facts must never choose + // the catalog used for settlement. + let billing_processing_tier = requested_processing_tier .as_deref() .map(canonical_processing_tier) - .or_else(|| { - requested_processing_tier.as_deref().map_or( - Some("standard".to_string()), - |requested| { - processing_tier_is_standard(requested).then(|| "standard".to_string()) - }, - ) - }); + .or_else(|| Some("standard".to_string())); let (tiered_pricing, tiered_pricing_source, processing_tier_price_multiplier) = match billing_processing_tier.as_deref() { @@ -208,34 +196,7 @@ impl BillingModelPricingSnapshot { return Ok(None); } - let mut billing_tiers = BTreeSet::from(["standard".to_string(), requested_billing_tier]); - for pricing in [ - self.model_tiered_pricing.as_ref(), - self.default_tiered_pricing.as_ref(), - ] - .into_iter() - .flatten() - { - let Some(processing_tiers) = pricing.get("processing_tiers").and_then(Value::as_object) - else { - continue; - }; - billing_tiers.extend(processing_tiers.keys().filter_map(|tier| { - normalize_processing_tier(Some(tier)).map(|tier| canonical_processing_tier(&tier)) - })); - } - - let mut candidates = Vec::new(); - for billing_tier in billing_tiers { - let resolution = self.authorization_pricing_for_tier( - requested_processing_tier.clone(), - Some(billing_tier), - )?; - if resolution.bills_standard_processing_tier() || resolution.tiered_pricing.is_some() { - candidates.push(resolution); - } - } - Ok((!candidates.is_empty()).then_some(candidates)) + Ok(Some(vec![requested_resolution])) } pub fn validate_authorization_pricing_configuration( @@ -839,7 +800,7 @@ mod tests { } #[test] - fn explicit_nonstandard_request_requires_actual_tier() { + fn explicit_nonstandard_request_selects_requested_catalog_without_actual_tier() { let pricing = snapshot( None, Some(json!({ @@ -852,17 +813,27 @@ mod tests { let resolution = pricing.resolve_pricing(Some("Priority"), None); - assert!(resolution.requires_actual_processing_tier()); + assert!(!resolution.requires_actual_processing_tier()); assert_eq!( resolution.requested_processing_tier.as_deref(), Some("priority") ); - assert_eq!(resolution.billing_processing_tier, None); - assert_eq!(resolution.tiered_pricing, None); + assert_eq!( + resolution.billing_processing_tier.as_deref(), + Some("priority") + ); + assert_eq!( + resolution + .tiered_pricing + .as_ref() + .and_then(|value| value.pointer("/tiers/0/input_price_per_1m")) + .and_then(serde_json::Value::as_f64), + Some(6.0) + ); } #[test] - fn actual_tier_selects_exact_catalog_and_source() { + fn response_actual_tier_is_audited_but_does_not_select_pricing_catalog() { let pricing = snapshot( Some(json!({ "processing_tiers": { @@ -878,28 +849,29 @@ mod tests { ); let flex = pricing.resolve_pricing(Some("priority"), Some("flex")); - assert_eq!(flex.billing_processing_tier.as_deref(), Some("flex")); + assert_eq!(flex.actual_processing_tier.as_deref(), Some("flex")); + assert_eq!(flex.billing_processing_tier.as_deref(), Some("priority")); assert_eq!( flex.tiered_pricing_source, - Some(BillingPricingSource::GlobalDefault) + Some(BillingPricingSource::ProviderOverride) ); assert_eq!( flex.tiered_pricing .as_ref() .and_then(|value| value.pointer("/tiers/0/input_price_per_1m")) .and_then(serde_json::Value::as_f64), - Some(1.5) + Some(9.0) ); let standard = pricing.resolve_pricing(Some("priority"), Some("Default")); assert_eq!(standard.actual_processing_tier.as_deref(), Some("default")); assert_eq!( standard.billing_processing_tier.as_deref(), - Some("standard") + Some("priority") ); assert_eq!( standard.tiered_pricing_source, - Some(BillingPricingSource::GlobalDefault) + Some(BillingPricingSource::ProviderOverride) ); } @@ -962,9 +934,7 @@ mod tests { .and_then(serde_json::Value::as_f64), Some(6.0) ); - assert!(candidates.iter().any(|resolution| { - resolution.billing_processing_tier.as_deref() == Some("standard") - })); + assert_eq!(candidates.len(), 1); } #[test] @@ -1222,7 +1192,7 @@ mod tests { } #[test] - fn actual_claude_fast_uses_the_exact_provider_multiplier_overlay() { + fn requested_claude_fast_uses_the_exact_provider_multiplier_overlay() { let pricing = snapshot( Some(json!({ "tiers": [{"up_to": null, "input_price_per_1m": 2.0}], @@ -1234,8 +1204,9 @@ mod tests { })), ); - let resolved = pricing.resolve_pricing(Some("priority"), Some("fast")); + let resolved = pricing.resolve_pricing(Some("fast"), Some("priority")); + assert_eq!(resolved.actual_processing_tier.as_deref(), Some("priority")); assert_eq!(resolved.billing_processing_tier.as_deref(), Some("fast")); assert_eq!( resolved diff --git a/crates/aether-billing/src/service.rs b/crates/aether-billing/src/service.rs index c19aba38e..d84169796 100644 --- a/crates/aether-billing/src/service.rs +++ b/crates/aether-billing/src/service.rs @@ -1104,7 +1104,7 @@ mod tests { } #[test] - fn nonstandard_request_without_actual_tier_fails_closed() { + fn nonstandard_request_without_actual_tier_uses_requested_catalog() { let result = BillingService::new() .calculate( &processing_pricing(), @@ -1112,57 +1112,49 @@ mod tests { ) .expect("billing should calculate"); - assert_eq!(result.cost_result.status, BillingSnapshotStatus::NoRule); - assert_eq!( - result.cost_result.snapshot.missing_required, - vec!["actual_processing_tier"] - ); + assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete); assert_eq!( result.cost_result.snapshot.resolved_dimensions["billing_processing_tier"], - json!(null) + json!("priority") + ); + assert_eq!( + result.cost_result.snapshot.resolved_variables["input_price_per_1m"], + json!(10.0) ); } #[test] - fn actual_tier_controls_standard_flex_and_priority_catalogs() { - let cases = [ - ( - "default", - 100, - 5.0, - 6.25, - BillingPricingSource::GlobalDefault, - ), - ("flex", 100, 2.5, 3.125, BillingPricingSource::GlobalDefault), - ( - "priority", - 100, - 10.0, - 12.5, - BillingPricingSource::ProviderOverride, - ), - ]; + fn response_actual_tier_does_not_override_requested_catalog() { + let cases = ["default", "flex", "priority"]; - for (actual, input_tokens, input_price, cache_write_price, source) in cases { + for actual in cases { let result = BillingService::new() .calculate( &processing_pricing(), - &processing_usage(Some("priority"), Some(actual), input_tokens), + &processing_usage(Some("priority"), Some(actual), 100), ) .expect("processing tier should resolve"); assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete); + assert_eq!( + result.pricing_resolution.actual_processing_tier.as_deref(), + Some(actual) + ); + assert_eq!( + result.pricing_resolution.billing_processing_tier.as_deref(), + Some("priority") + ); assert_eq!( result.cost_result.snapshot.resolved_variables["input_price_per_1m"], - json!(input_price) + json!(10.0) ); assert_eq!( result.cost_result.snapshot.resolved_variables["cache_creation_price_per_1m"], - json!(cache_write_price) + json!(12.5) ); assert_eq!( result.pricing_resolution.tiered_pricing_source, - Some(source) + Some(BillingPricingSource::ProviderOverride) ); } } @@ -1373,7 +1365,7 @@ mod tests { } #[test] - fn finite_processing_catalog_and_unknown_actual_tier_fail_closed() { + fn finite_processing_catalog_fails_only_on_requested_catalog_bounds() { let priority = BillingService::new() .calculate( &processing_pricing(), @@ -1386,16 +1378,22 @@ mod tests { vec!["input_context_tier"] ); - let unknown = BillingService::new() + let conflicting_actual = BillingService::new() .calculate( &processing_pricing(), &processing_usage(Some("priority"), Some("expedited"), 100), ) .expect("billing should calculate"); - assert_eq!(unknown.cost_result.status, BillingSnapshotStatus::NoRule); assert_eq!( - unknown.cost_result.snapshot.missing_required, - vec!["processing_tier_catalog"] + conflicting_actual.cost_result.status, + BillingSnapshotStatus::Complete + ); + assert_eq!( + conflicting_actual + .pricing_resolution + .billing_processing_tier + .as_deref(), + Some("priority") ); } @@ -1502,13 +1500,17 @@ mod tests { estimate.max_output_tokens = Some(0); estimate.cache_ttl_minutes = Some(30); - for requested_processing_tier in [None, Some("standard"), Some("flex")] { + for (requested_processing_tier, expected) in [ + (None, 3.75), + (Some("standard"), 3.75), + (Some("flex"), 1.875), + ] { estimate.requested_processing_tier = requested_processing_tier.map(ToOwned::to_owned); assert_eq!( service .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) .expect("eligible processing catalogs should calculate"), - Some(3.75), + Some(expected), "requested tier: {requested_processing_tier:?}" ); } @@ -1523,7 +1525,7 @@ mod tests { } #[test] - fn unknown_actual_tier_cannot_fall_back_to_fixed_request_price() { + fn unknown_actual_tier_does_not_override_requested_tier_or_fixed_price() { let pricing = BillingModelPricingSnapshot { default_price_per_request: Some(0.02), ..processing_pricing() @@ -1535,16 +1537,20 @@ mod tests { ) .expect("billing should calculate"); - assert_eq!(result.cost_result.status, BillingSnapshotStatus::NoRule); + assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete); assert_eq!( - result.cost_result.snapshot.missing_required, - vec!["processing_tier_catalog"] + result.pricing_resolution.billing_processing_tier.as_deref(), + Some("priority") + ); + assert_eq!( + result.pricing_resolution.actual_processing_tier.as_deref(), + Some("expedited") ); assert_eq!(result.pricing_resolution.price_per_request, Some(0.02)); } #[test] - fn authorization_estimate_bounds_requested_and_provider_actual_catalogs() { + fn authorization_estimate_bounds_only_the_requested_catalog() { let service = BillingService::new(); let mut estimate = BillingAuthorizationEstimateInput::new("chat", 100_000); estimate.api_format = Some("openai:responses".to_string()); @@ -1563,8 +1569,8 @@ mod tests { .expect("flex estimate should be bounded"); assert_eq!(priority, 61.25); - assert_eq!(flex, 61.25); - assert_eq!(priority, flex); + assert_eq!(flex, 15.3125); + assert!(priority > flex); } #[test] diff --git a/crates/aether-data/adapters/mysql/src/usage.rs b/crates/aether-data/adapters/mysql/src/usage.rs index 3c57a9379..b8ca918f6 100644 --- a/crates/aether-data/adapters/mysql/src/usage.rs +++ b/crates/aether-data/adapters/mysql/src/usage.rs @@ -135,22 +135,22 @@ INSERT INTO `usage` ( ON DUPLICATE KEY UPDATE user_id = VALUES(user_id), api_key_id = VALUES(api_key_id), - provider_name = VALUES(provider_name), - model = VALUES(model), - target_model = VALUES(target_model), - provider_id = VALUES(provider_id), - provider_endpoint_id = VALUES(provider_endpoint_id), - provider_api_key_id = VALUES(provider_api_key_id), - request_type = VALUES(request_type), - api_format = VALUES(api_format), - api_family = VALUES(api_family), - endpoint_kind = VALUES(endpoint_kind), - endpoint_api_format = VALUES(endpoint_api_format), - provider_api_family = VALUES(provider_api_family), - provider_endpoint_kind = VALUES(provider_endpoint_kind), - has_format_conversion = VALUES(has_format_conversion), - is_stream = VALUES(is_stream), - upstream_is_stream = VALUES(upstream_is_stream), + provider_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_name ELSE VALUES(provider_name) END, + model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN model ELSE VALUES(model) END, + target_model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN target_model ELSE VALUES(target_model) END, + provider_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_id ELSE VALUES(provider_id) END, + provider_endpoint_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_id ELSE VALUES(provider_endpoint_id) END, + provider_api_key_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_key_id ELSE VALUES(provider_api_key_id) END, + request_type = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_type ELSE VALUES(request_type) END, + api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_format ELSE VALUES(api_format) END, + api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_family ELSE VALUES(api_family) END, + endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_kind ELSE VALUES(endpoint_kind) END, + endpoint_api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_api_format ELSE VALUES(endpoint_api_format) END, + provider_api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_family ELSE VALUES(provider_api_family) END, + provider_endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_kind ELSE VALUES(provider_endpoint_kind) END, + has_format_conversion = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN has_format_conversion ELSE VALUES(has_format_conversion) END, + is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN is_stream ELSE VALUES(is_stream) END, + upstream_is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN upstream_is_stream ELSE VALUES(upstream_is_stream) END, input_tokens = CASE WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN input_tokens ELSE VALUES(input_tokens) @@ -231,15 +231,15 @@ ON DUPLICATE KEY UPDATE WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN billing_status ELSE VALUES(billing_status) END, - request_metadata = VALUES(request_metadata), - candidate_id = VALUES(candidate_id), - candidate_index = VALUES(candidate_index), - key_name = VALUES(key_name), - planner_kind = VALUES(planner_kind), - route_family = VALUES(route_family), - route_kind = VALUES(route_kind), - execution_path = VALUES(execution_path), - local_execution_runtime_miss_reason = VALUES(local_execution_runtime_miss_reason), + request_metadata = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_metadata ELSE VALUES(request_metadata) END, + candidate_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_id ELSE VALUES(candidate_id) END, + candidate_index = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_index ELSE VALUES(candidate_index) END, + key_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN key_name ELSE VALUES(key_name) END, + planner_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN planner_kind ELSE VALUES(planner_kind) END, + route_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_family ELSE VALUES(route_family) END, + route_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_kind ELSE VALUES(route_kind) END, + execution_path = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN execution_path ELSE VALUES(execution_path) END, + local_execution_runtime_miss_reason = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN local_execution_runtime_miss_reason ELSE VALUES(local_execution_runtime_miss_reason) END, finalized_at = CASE WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN finalized_at ELSE VALUES(finalized_at) diff --git a/crates/aether-data/adapters/mysql/src/usage/tests.rs b/crates/aether-data/adapters/mysql/src/usage/tests.rs index ebfe9027d..9d60fd604 100644 --- a/crates/aether-data/adapters/mysql/src/usage/tests.rs +++ b/crates/aether-data/adapters/mysql/src/usage/tests.rs @@ -65,6 +65,50 @@ fn mysql_usage_upsert_keeps_terminal_state_when_streaming_arrives_late() { )); } +#[test] +fn mysql_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() { + for field in [ + "provider_name", + "model", + "target_model", + "provider_id", + "provider_endpoint_id", + "provider_api_key_id", + "request_type", + "api_format", + "api_family", + "endpoint_kind", + "endpoint_api_format", + "provider_api_family", + "provider_endpoint_kind", + "has_format_conversion", + "is_stream", + "upstream_is_stream", + "request_metadata", + "candidate_id", + "candidate_index", + "key_name", + "planner_kind", + "route_family", + "route_kind", + "execution_path", + "local_execution_runtime_miss_reason", + ] { + let assignment = format!("{field} = CASE WHEN ("); + assert!( + super::UPSERT_USAGE_SQL.contains(&assignment), + "missing lifecycle guard for {field}" + ); + let preserve = format!("THEN {field} ELSE VALUES({field}) END"); + assert!( + super::UPSERT_USAGE_SQL.contains(&preserve), + "late lifecycle must preserve {field}" + ); + } + assert!(super::UPSERT_USAGE_SQL + .contains("OR (status = 'streaming' AND VALUES(status) = 'pending')")); +} + #[tokio::test] async fn mysql_usage_write_repository_upserts_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") diff --git a/crates/aether-data/adapters/postgres/src/usage/mod.rs b/crates/aether-data/adapters/postgres/src/usage/mod.rs index 5848f29d8..a1bfcd4a1 100644 --- a/crates/aether-data/adapters/postgres/src/usage/mod.rs +++ b/crates/aether-data/adapters/postgres/src/usage/mod.rs @@ -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, request_metadata_json: Option, + 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, +) -> 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 { @@ -10485,40 +10655,78 @@ fn json_bind_text(value: Option<&Value>) -> Result, DataLayerErro fn prepare_usage_upsert_context( usage: &UpsertUsageRecord, ) -> Result { + 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()?; diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql b/crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql index 4859983eb..0d3f2afed 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql @@ -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', diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql b/crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql index 4859983eb..0d3f2afed 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql @@ -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', diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_sql.sql b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_sql.sql index 4c616d3c8..3acbc2680 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_sql.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_sql.sql @@ -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( diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_http_audit_sql.sql b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_http_audit_sql.sql index 87e084c91..0779098c6 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_http_audit_sql.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_http_audit_sql.sql @@ -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() diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_routing_snapshot_sql.sql b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_routing_snapshot_sql.sql index 6bfa32719..636a78f87 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_routing_snapshot_sql.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_routing_snapshot_sql.sql @@ -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() diff --git a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_settlement_pricing_snapshot_sql.sql b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_settlement_pricing_snapshot_sql.sql index 73cd637ca..8a69b8fa5 100644 --- a/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_settlement_pricing_snapshot_sql.sql +++ b/crates/aether-data/adapters/postgres/src/usage/queries/upsert_usage_settlement_pricing_snapshot_sql.sql @@ -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() diff --git a/crates/aether-data/adapters/postgres/src/usage/tests.rs b/crates/aether-data/adapters/postgres/src/usage/tests.rs index 93dc30d2d..eeeb6bc62 100644 --- a/crates/aether-data/adapters/postgres/src/usage/tests.rs +++ b/crates/aether-data/adapters/postgres/src/usage/tests.rs @@ -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::, _>("settlement_snapshot") + .expect("settlement snapshot should decode") + .is_none()); + assert!(settlement_row + .try_get::, _>("billing_dimensions") + .expect("billing dimensions should decode") + .is_none()); + assert!(settlement_row + .try_get::, _>("rate_multiplier") + .expect("rate multiplier should decode") + .is_none()); + assert!(settlement_row + .try_get::, _>("input_price_per_1m") + .expect("input price should decode") + .is_none()); + assert!(settlement_row + .try_get::, _>("output_price_per_1m") + .expect("output price should decode") + .is_none()); + + let physical_metadata = sqlx::query_scalar::<_, Option>( + "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::, _>("provider_request_body") + .expect("provider body column should decode") + .is_none()); + assert!(physical_body + .try_get::>, _>("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::, _>("provider_request_body_ref") + .expect("provider body ref should decode") + .is_none()); + assert_eq!( + physical_http + .try_get::, _>("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!({ diff --git a/crates/aether-data/adapters/sqlite/src/usage.rs b/crates/aether-data/adapters/sqlite/src/usage.rs index a328636ae..49e0930d3 100644 --- a/crates/aether-data/adapters/sqlite/src/usage.rs +++ b/crates/aether-data/adapters/sqlite/src/usage.rs @@ -156,22 +156,22 @@ INSERT INTO "usage" ( ON CONFLICT (request_id) DO UPDATE SET user_id = excluded.user_id, api_key_id = excluded.api_key_id, - provider_name = excluded.provider_name, - model = excluded.model, - target_model = excluded.target_model, - provider_id = excluded.provider_id, - provider_endpoint_id = excluded.provider_endpoint_id, - provider_api_key_id = excluded.provider_api_key_id, - request_type = excluded.request_type, - api_format = excluded.api_format, - api_family = excluded.api_family, - endpoint_kind = excluded.endpoint_kind, - endpoint_api_format = excluded.endpoint_api_format, - provider_api_family = excluded.provider_api_family, - provider_endpoint_kind = excluded.provider_endpoint_kind, - has_format_conversion = excluded.has_format_conversion, - is_stream = excluded.is_stream, - upstream_is_stream = excluded.upstream_is_stream, + provider_name = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_name ELSE excluded.provider_name END, + model = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".model ELSE excluded.model END, + target_model = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".target_model ELSE excluded.target_model END, + provider_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_id ELSE excluded.provider_id END, + provider_endpoint_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_endpoint_id ELSE excluded.provider_endpoint_id END, + provider_api_key_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_api_key_id ELSE excluded.provider_api_key_id END, + request_type = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".request_type ELSE excluded.request_type END, + api_format = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".api_format ELSE excluded.api_format END, + api_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".api_family ELSE excluded.api_family END, + endpoint_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".endpoint_kind ELSE excluded.endpoint_kind END, + endpoint_api_format = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".endpoint_api_format ELSE excluded.endpoint_api_format END, + provider_api_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_api_family ELSE excluded.provider_api_family END, + provider_endpoint_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".provider_endpoint_kind ELSE excluded.provider_endpoint_kind END, + has_format_conversion = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".has_format_conversion ELSE excluded.has_format_conversion END, + is_stream = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".is_stream ELSE excluded.is_stream END, + upstream_is_stream = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".upstream_is_stream ELSE excluded.upstream_is_stream END, input_tokens = CASE WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".input_tokens ELSE excluded.input_tokens @@ -255,15 +255,15 @@ ON CONFLICT (request_id) DO UPDATE SET WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".billing_status ELSE excluded.billing_status END, - request_metadata = excluded.request_metadata, - candidate_id = COALESCE(excluded.candidate_id, "usage".candidate_id), - candidate_index = COALESCE(excluded.candidate_index, "usage".candidate_index), - key_name = COALESCE(excluded.key_name, "usage".key_name), - planner_kind = excluded.planner_kind, - route_family = excluded.route_family, - route_kind = excluded.route_kind, - execution_path = excluded.execution_path, - local_execution_runtime_miss_reason = excluded.local_execution_runtime_miss_reason, + request_metadata = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".request_metadata ELSE excluded.request_metadata END, + candidate_id = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".candidate_id WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.candidate_id ELSE COALESCE(excluded.candidate_id, "usage".candidate_id) END, + candidate_index = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".candidate_index WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.candidate_index ELSE COALESCE(excluded.candidate_index, "usage".candidate_index) END, + key_name = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".key_name WHEN excluded.status IN ('completed', 'failed', 'cancelled') THEN excluded.key_name ELSE COALESCE(excluded.key_name, "usage".key_name) END, + planner_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".planner_kind ELSE excluded.planner_kind END, + route_family = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".route_family ELSE excluded.route_family END, + route_kind = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".route_kind ELSE excluded.route_kind END, + execution_path = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".execution_path ELSE excluded.execution_path END, + local_execution_runtime_miss_reason = CASE WHEN ("usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming')) OR ("usage".status = 'streaming' AND excluded.status = 'pending') THEN "usage".local_execution_runtime_miss_reason ELSE excluded.local_execution_runtime_miss_reason END, finalized_at = CASE WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".finalized_at ELSE excluded.finalized_at diff --git a/crates/aether-data/adapters/sqlite/src/usage/tests.rs b/crates/aether-data/adapters/sqlite/src/usage/tests.rs index 527848e43..aebfb0c2d 100644 --- a/crates/aether-data/adapters/sqlite/src/usage/tests.rs +++ b/crates/aether-data/adapters/sqlite/src/usage/tests.rs @@ -6,6 +6,49 @@ use aether_data_contracts::repository::usage::{ UsageReadRepository, UsageTimeSeriesGranularity, UsageWriteRepository, }; +#[test] +fn sqlite_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_lifecycle() { + for field in [ + "provider_name", + "model", + "target_model", + "provider_id", + "provider_endpoint_id", + "provider_api_key_id", + "request_type", + "api_format", + "api_family", + "endpoint_kind", + "endpoint_api_format", + "provider_api_family", + "provider_endpoint_kind", + "has_format_conversion", + "is_stream", + "upstream_is_stream", + "request_metadata", + "candidate_id", + "candidate_index", + "key_name", + "planner_kind", + "route_family", + "route_kind", + "execution_path", + "local_execution_runtime_miss_reason", + ] { + let assignment = format!("{field} = CASE WHEN ("); + assert!( + super::UPSERT_USAGE_SQL.contains(&assignment), + "missing lifecycle guard for {field}" + ); + assert!( + super::UPSERT_USAGE_SQL.contains(&format!("THEN \"usage\".{field}")), + "late lifecycle must preserve {field}" + ); + } + assert!(super::UPSERT_USAGE_SQL + .contains("OR (\"usage\".status = 'streaming' AND excluded.status = 'pending')")); +} + #[tokio::test] async fn sqlite_provider_performance_can_skip_timeline() { let pool = sqlx::sqlite::SqlitePoolOptions::new() @@ -139,7 +182,7 @@ async fn sqlite_usage_write_repository_does_not_regress_void_usage() { .expect("sqlite migrations should run"); seed_stats_targets(&pool).await; - let repository = SqliteUsageWriteRepository::new(pool); + let repository = SqliteUsageWriteRepository::new(pool.clone()); repository .upsert(sample_usage("request-1", "failed", "void", 1_000)) .await @@ -205,9 +248,11 @@ async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late .expect("sqlite migrations should run"); seed_stats_targets(&pool).await; - let repository = SqliteUsageWriteRepository::new(pool); + let repository = SqliteUsageWriteRepository::new(pool.clone()); + let mut terminal = sample_usage("request-1", "completed", "pending", 1_000); + terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal-trace"})); repository - .upsert(sample_usage("request-1", "completed", "pending", 1_000)) + .upsert(terminal) .await .expect("terminal usage should upsert"); @@ -222,6 +267,30 @@ async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late late_streaming.response_time_ms = Some(9_999); late_streaming.first_byte_time_ms = Some(9_999); late_streaming.finalized_at_unix_secs = None; + late_streaming.provider_name = "Late Provider".to_string(); + late_streaming.model = "late-model".to_string(); + late_streaming.target_model = Some("late-target".to_string()); + late_streaming.request_type = Some("late-request".to_string()); + late_streaming.api_format = Some("late:api".to_string()); + late_streaming.api_family = Some("late-family".to_string()); + late_streaming.endpoint_kind = Some("late-endpoint".to_string()); + late_streaming.endpoint_api_format = Some("late:endpoint".to_string()); + late_streaming.provider_api_family = Some("late-provider-family".to_string()); + late_streaming.provider_endpoint_kind = Some("late-provider-endpoint".to_string()); + late_streaming.has_format_conversion = Some(false); + late_streaming.is_stream = Some(true); + late_streaming.candidate_id = Some("late-candidate".to_string()); + late_streaming.candidate_index = Some(99); + late_streaming.key_name = Some("late-key".to_string()); + late_streaming.planner_kind = Some("late-planner".to_string()); + late_streaming.route_family = Some("late-route-family".to_string()); + late_streaming.route_kind = Some("late-route-kind".to_string()); + late_streaming.execution_path = Some("late-path".to_string()); + late_streaming.local_execution_runtime_miss_reason = Some("late-miss".to_string()); + late_streaming.request_metadata = Some(serde_json::json!({ + "provider_service_tier": "priority", + "upstream_is_stream": true + })); let current = repository .upsert(late_streaming) @@ -238,6 +307,91 @@ async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late assert_eq!(current.first_byte_time_ms, Some(12)); assert_eq!(current.finalized_at_unix_secs, Some(1_000)); assert_eq!(current.updated_at_unix_secs, 1_000); + assert_eq!(current.provider_name, "Provider One"); + assert_eq!(current.model, "model-1"); + assert_eq!(current.target_model.as_deref(), Some("target-model")); + assert_eq!(current.request_type.as_deref(), Some("chat")); + assert_eq!(current.api_format.as_deref(), Some("openai")); + assert!(current.has_format_conversion); + assert!(!current.is_stream); + assert_eq!(current.candidate_id.as_deref(), Some("candidate-1")); + assert_eq!(current.candidate_index, Some(1)); + assert_eq!(current.key_name.as_deref(), Some("key-one")); + assert_eq!(current.planner_kind.as_deref(), Some("default")); + assert_eq!(current.route_family.as_deref(), Some("chat")); + assert_eq!(current.route_kind.as_deref(), Some("completion")); + assert_eq!(current.execution_path.as_deref(), Some("remote")); + assert_eq!(current.provider_service_tier(), None); + assert_eq!( + current + .request_metadata + .as_ref() + .and_then(|value| value.get("trace_id")) + .and_then(serde_json::Value::as_str), + Some("terminal-trace") + ); + + let listed = SqliteUsageReadRepository::new(pool) + .list_usage_audits(&UsageAuditListQuery { + limit: Some(10), + newest_first: true, + ..UsageAuditListQuery::default() + }) + .await + .expect("usage list should load") + .into_iter() + .find(|item| item.request_id == "request-1") + .expect("terminal usage should be listed"); + assert_eq!(listed.provider_name, "Provider One"); + assert_eq!(listed.model, "model-1"); + assert_eq!(listed.candidate_id.as_deref(), Some("candidate-1")); + assert_eq!(listed.provider_service_tier(), None); +} + +#[tokio::test] +async fn sqlite_usage_write_repository_allows_authoritative_completed_recovery() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let repository = SqliteUsageWriteRepository::new(pool); + repository + .upsert(sample_usage("request-recovery", "failed", "void", 1_000)) + .await + .expect("void failure should upsert"); + + let mut recovery = sample_usage("request-recovery", "completed", "pending", 1_001); + recovery.provider_name = "Recovered Provider".to_string(); + recovery.model = "recovered-model".to_string(); + recovery.target_model = Some("recovered-target".to_string()); + recovery.api_format = Some("recovered:api".to_string()); + recovery.candidate_id = Some("recovered-candidate".to_string()); + recovery.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"})); + let recovered = repository + .upsert(recovery) + .await + .expect("completed recovery should upsert"); + + assert_eq!(recovered.status, "completed"); + assert_eq!(recovered.billing_status, "pending"); + assert_eq!(recovered.provider_name, "Recovered Provider"); + assert_eq!(recovered.model, "recovered-model"); + assert_eq!(recovered.target_model.as_deref(), Some("recovered-target")); + assert_eq!(recovered.api_format.as_deref(), Some("recovered:api")); + assert_eq!( + recovered.candidate_id.as_deref(), + Some("recovered-candidate") + ); + assert_eq!( + recovered.provider_service_tier().as_deref(), + Some("priority") + ); } #[tokio::test] @@ -279,6 +433,51 @@ async fn sqlite_usage_write_repository_preserves_streaming_response_start_from_l assert_eq!(current.first_byte_time_ms, Some(12)); } +#[tokio::test] +async fn sqlite_usage_write_repository_keeps_streaming_capture_from_late_pending() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let repository = SqliteUsageWriteRepository::new(pool); + let mut streaming = sample_usage("request-streaming-capture", "streaming", "pending", 1_000); + streaming.request_metadata = Some(serde_json::json!({"trace_id": "streaming-final"})); + repository + .upsert(streaming) + .await + .expect("streaming usage should upsert"); + + let mut late_pending = sample_usage("request-streaming-capture", "pending", "pending", 1_001); + late_pending.provider_name = "Late Provider".to_string(); + late_pending.model = "late-model".to_string(); + late_pending.candidate_id = Some("late-candidate".to_string()); + late_pending.request_metadata = Some(serde_json::json!({"provider_service_tier": "priority"})); + let current = repository + .upsert(late_pending) + .await + .expect("late pending usage should not regress streaming capture"); + + assert_eq!(current.status, "streaming"); + assert_eq!(current.provider_name, "Provider One"); + assert_eq!(current.model, "model-1"); + assert_eq!(current.candidate_id.as_deref(), Some("candidate-1")); + assert_eq!(current.provider_service_tier(), None); + assert_eq!( + current + .request_metadata + .as_ref() + .and_then(|value| value.get("trace_id")) + .and_then(serde_json::Value::as_str), + Some("streaming-final") + ); +} + #[tokio::test] async fn sqlite_usage_write_repository_cleans_stale_pending_requests() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/contracts/src/repository/usage/mod.rs b/crates/aether-data/contracts/src/repository/usage/mod.rs index 30be92f09..864ed34ce 100644 --- a/crates/aether-data/contracts/src/repository/usage/mod.rs +++ b/crates/aether-data/contracts/src/repository/usage/mod.rs @@ -6,9 +6,10 @@ pub use types::{ extract_provider_actual_service_tier_from_response, extract_provider_cache_ttl_minutes_from_metadata, extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body, normalize_provider_service_tier, parse_usage_body_ref, - resolve_provider_cache_ttl_minutes, usage_body_ref, usage_request_metadata_client_family, - ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary, - ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary, + resolve_provider_cache_ttl_minutes, resolve_provider_service_tier_from_request_capture, + usage_body_ref, usage_request_metadata_client_family, ApiKeyLastUsedDelta, + ManagementTokenCounterDelta, PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, + ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary, @@ -34,5 +35,5 @@ pub use types::{ UsageReadRepository, UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery, UsageWriteRepository, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, - PROVIDER_SERVICE_TIER_METADATA_KEY, + PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, }; diff --git a/crates/aether-data/contracts/src/repository/usage/types.rs b/crates/aether-data/contracts/src/repository/usage/types.rs index da9b02fe0..810314401 100644 --- a/crates/aether-data/contracts/src/repository/usage/types.rs +++ b/crates/aether-data/contracts/src/repository/usage/types.rs @@ -3,6 +3,7 @@ use chrono::{DateTime, Utc}; use serde_json::Value; pub const PROVIDER_REASONING_EFFORT_METADATA_KEY: &str = "provider_reasoning_effort"; +pub const REQUESTED_REASONING_EFFORT_METADATA_KEY: &str = "requested_reasoning_effort"; pub const PROVIDER_SERVICE_TIER_METADATA_KEY: &str = "provider_service_tier"; pub const PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY: &str = "provider_actual_service_tier"; pub const PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY: &str = "provider_cache_ttl_minutes"; @@ -95,6 +96,89 @@ pub fn normalize_provider_service_tier(value: &str) -> Option { Some(normalized) } +/// Resolves a provider processing tier exclusively from the final upstream request. +/// +/// A complete captured body is authoritative, including when it contains no tier. The metadata +/// fallback is reserved for bodies that were stripped, externalized, or truncated after the tier +/// had already been derived from that same request. +pub fn resolve_provider_service_tier_from_request_capture( + provider_request_body: Option<&Value>, + provider_request_body_state: Option, + request_metadata: Option<&Value>, +) -> Option { + if request_body_capture_is_authoritative(provider_request_body, provider_request_body_state) { + return extract_provider_service_tier_from_body(provider_request_body); + } + + if !matches!( + provider_request_body_state, + Some( + UsageBodyCaptureState::Inline + | UsageBodyCaptureState::Reference + | UsageBodyCaptureState::Truncated + | UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable + ) + ) { + return None; + } + + request_metadata + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(PROVIDER_SERVICE_TIER_METADATA_KEY)) + .and_then(Value::as_str) + .and_then(normalize_provider_service_tier) +} + +fn request_body_capture_is_authoritative( + request_body: Option<&Value>, + request_body_state: Option, +) -> bool { + let Some(request_body) = request_body else { + return false; + }; + if matches!( + request_body_state, + Some( + UsageBodyCaptureState::None + | UsageBodyCaptureState::Truncated + | UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable + ) + ) { + 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 resolve_reasoning_effort_from_request_capture( + request_body: Option<&Value>, + request_body_state: Option, + request_metadata: Option<&Value>, + metadata_key: &str, +) -> Option { + if request_body_capture_is_authoritative(request_body, request_body_state) { + return extract_provider_reasoning_effort_from_body(request_body); + } + + // An explicit `none` marker describes the final capture attempt and must win over any stale + // inline body/metadata left by an earlier candidate. `None` (the absence of a marker) remains + // the legacy list-query shape, where metadata is the only available representation. + if request_body_state == Some(UsageBodyCaptureState::None) { + return None; + } + + request_metadata + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(metadata_key)) + .and_then(Value::as_str) + .and_then(normalize_provider_reasoning_effort) +} + pub fn resolve_provider_cache_ttl_minutes( provider_api_format: Option<&str>, provider_model: Option<&str>, @@ -511,33 +595,40 @@ impl StoredRequestUsageAudit { } pub fn provider_reasoning_effort(&self) -> Option { - if self - .provider_request_body - .as_ref() - .and_then(Value::as_object) - .is_some() - { - return extract_provider_reasoning_effort_from_body( - self.provider_request_body.as_ref(), - ); - } + resolve_reasoning_effort_from_request_capture( + self.provider_request_body.as_ref(), + self.provider_request_body_state, + self.request_metadata.as_ref(), + PROVIDER_REASONING_EFFORT_METADATA_KEY, + ) + } - self.request_metadata_string(PROVIDER_REASONING_EFFORT_METADATA_KEY) - .and_then(normalize_provider_reasoning_effort) + pub fn requested_reasoning_effort(&self) -> Option { + resolve_reasoning_effort_from_request_capture( + self.request_body.as_ref(), + self.request_body_state, + self.request_metadata.as_ref(), + REQUESTED_REASONING_EFFORT_METADATA_KEY, + ) } pub fn provider_service_tier(&self) -> Option { - if self - .provider_request_body - .as_ref() - .and_then(Value::as_object) - .is_some() - { - return extract_provider_service_tier_from_body(self.provider_request_body.as_ref()); - } - - self.request_metadata_string(PROVIDER_SERVICE_TIER_METADATA_KEY) - .and_then(normalize_provider_service_tier) + resolve_provider_service_tier_from_request_capture( + self.provider_request_body.as_ref(), + self.provider_request_body_state, + self.request_metadata.as_ref(), + ) + .or_else(|| { + // Lightweight list queries intentionally omit bodies and typed capture state. Their + // request metadata was normalized before persistence and is the only available copy + // of the final provider-request fact. + if self.provider_request_body.is_none() && self.provider_request_body_state.is_none() { + self.request_metadata_string(PROVIDER_SERVICE_TIER_METADATA_KEY) + .and_then(normalize_provider_service_tier) + } else { + None + } + }) } pub fn provider_actual_service_tier(&self) -> Option { @@ -2684,6 +2775,126 @@ mod tests { assert_eq!(usage.provider_service_tier(), None); } + #[test] + fn requested_and_provider_reasoning_efforts_remain_independent() { + let mut usage = sample_usage(); + usage.request_body = Some(json!({ + "reasoning": { "effort": "XHigh" } + })); + usage.provider_request_body = Some(json!({ + "reasoning_effort": "max" + })); + usage.request_metadata = Some(json!({ + "requested_reasoning_effort": "low", + "provider_reasoning_effort": "medium" + })); + + assert_eq!(usage.requested_reasoning_effort().as_deref(), Some("xhigh")); + assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max")); + + usage.request_body = Some(json!({ "model": "gpt-5" })); + assert_eq!(usage.requested_reasoning_effort(), None); + + usage.request_body = None; + assert_eq!(usage.requested_reasoning_effort().as_deref(), Some("low")); + } + + #[test] + fn explicit_none_capture_ignores_stale_request_bodies_and_metadata() { + let mut usage = sample_usage(); + usage.request_body = Some(json!({ + "reasoning": { "effort": "xhigh" } + })); + usage.request_body_state = Some(UsageBodyCaptureState::None); + usage.provider_request_body = Some(json!({ + "reasoning_effort": "max", + "service_tier": "priority" + })); + usage.provider_request_body_state = Some(UsageBodyCaptureState::None); + usage.request_metadata = Some(json!({ + "requested_reasoning_effort": "high", + "provider_reasoning_effort": "medium", + "provider_service_tier": "priority" + })); + + assert_eq!(usage.requested_reasoning_effort(), None); + assert_eq!(usage.provider_reasoning_effort(), None); + assert_eq!(usage.provider_service_tier(), None); + } + + #[test] + fn reasoning_mapping_uses_derived_metadata_after_request_bodies_are_truncated() { + let truncated = json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded" + }); + let mut usage = sample_usage(); + usage.request_body = Some(truncated.clone()); + usage.request_body_state = Some(UsageBodyCaptureState::Truncated); + usage.provider_request_body = Some(truncated); + usage.provider_request_body_state = Some(UsageBodyCaptureState::Truncated); + usage.request_metadata = Some(json!({ + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max" + })); + + assert_eq!(usage.requested_reasoning_effort().as_deref(), Some("xhigh")); + assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max")); + } + + #[test] + fn provider_service_tier_does_not_infer_fast_from_response_or_stale_metadata() { + let mut usage = sample_usage(); + usage.provider_request_body = Some(json!({ + "model": "gpt-5" + })); + usage.request_metadata = Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "priority" + })); + usage.response_body = Some(json!({ + "service_tier": "priority" + })); + + assert_eq!(usage.provider_service_tier(), None); + } + + #[test] + fn provider_service_tier_uses_derived_metadata_after_final_body_is_stripped() { + let mut usage = sample_usage(); + usage.provider_request_body = None; + usage.provider_request_body_state = Some(UsageBodyCaptureState::Disabled); + usage.request_metadata = Some(json!({ + "provider_service_tier": "priority" + })); + usage.response_body = Some(json!({ + "service_tier": "default" + })); + + assert_eq!(usage.provider_service_tier().as_deref(), Some("priority")); + } + + #[test] + fn provider_service_tier_uses_request_derived_metadata_after_body_truncation() { + let mut usage = sample_usage(); + usage.provider_request_body = Some(json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded", + "max_bytes": 128, + "source_bytes": 4096, + "value_kind": "object" + })); + usage.provider_request_body_state = Some(UsageBodyCaptureState::Truncated); + usage.request_metadata = Some(json!({ + "provider_service_tier": "priority" + })); + usage.response_body = Some(json!({ + "service_tier": "default" + })); + + assert_eq!(usage.provider_service_tier().as_deref(), Some("priority")); + } + #[test] fn provider_cache_ttl_uses_final_openai_contract_then_preserved_metadata() { let mut usage = sample_usage(); diff --git a/crates/aether-data/runtime/src/repository/usage/memory.rs b/crates/aether-data/runtime/src/repository/usage/memory.rs index 915fbcb2c..ed1542356 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory.rs @@ -14,14 +14,16 @@ use aether_data_contracts::repository::usage::{ StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary, StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, - UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, + UsageBodyCaptureState, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageSettledCostSummaryQuery, - UsageTimeSeriesGranularity, UsageTimeSeriesQuery, + UsageTimeSeriesGranularity, UsageTimeSeriesQuery, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, + PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, + REQUESTED_REASONING_EFFORT_METADATA_KEY, }; use async_trait::async_trait; use chrono::Utc; @@ -2703,10 +2705,14 @@ fn hydrate_client_family(item: &mut StoredRequestUsageAudit) { fn persisted_usage_body_ref( incoming_ref: Option<&str>, incoming_body: Option<&Value>, + incoming_state: Option, _metadata: Option<&Value>, existing: Option<&StoredRequestUsageAudit>, field: UsageBodyField, ) -> Option { + if incoming_state == Some(UsageBodyCaptureState::None) { + return None; + } if incoming_body.is_some() { return None; } @@ -2724,6 +2730,78 @@ fn persisted_usage_body_ref( }) } +fn request_body_capture_replaces_derived_facts( + request_body: Option<&Value>, + request_body_state: Option, +) -> bool { + 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_request_body_facts( + metadata: Option<&Value>, + clear_client_request_body_facts: bool, + clear_provider_request_body_facts: bool, +) -> Value { + let mut metadata = 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"); + } + Value::Object(metadata) +} + +fn retain_previous_request_audit_metadata( + metadata: Option<&Value>, + preserve_client_request_body_facts: bool, +) -> Value { + let Some(metadata) = metadata.and_then(Value::as_object) else { + return Value::Object(serde_json::Map::new()); + }; + let mut retained = serde_json::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) = 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) = metadata.get(key) { + retained.insert(key.to_string(), value.clone()); + } + } + } + Value::Object(retained) +} + fn merge_usage_status_code( existing: Option<&StoredRequestUsageAudit>, incoming_status: &str, @@ -2789,14 +2867,57 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { return Ok(existing.expect("existing usage should be present").clone()); } - let request_metadata = usage.request_metadata.clone().or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.request_metadata.clone()) + 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); + let replace_routing_snapshot = usage_status_is_finalized(&usage.status); + let mut incoming_request_metadata = usage.request_metadata.clone(); + if incoming_request_metadata.is_some() + && (clear_request_body || clear_provider_request_body) + { + incoming_request_metadata = Some(clear_request_body_facts( + incoming_request_metadata.as_ref(), + clear_request_body, + clear_provider_request_body, + )); + } + let request_metadata = incoming_request_metadata.or_else(|| { + if replace_routing_snapshot { + Some(retain_previous_request_audit_metadata( + existing + .as_ref() + .and_then(|existing| existing.request_metadata.as_ref()), + !replace_client_request_body_facts, + )) + } else if replace_client_request_body_facts || replace_provider_request_body_facts { + Some(clear_request_body_facts( + existing + .as_ref() + .and_then(|existing| existing.request_metadata.as_ref()), + replace_client_request_body_facts, + replace_provider_request_body_facts, + )) + } else { + existing + .as_ref() + .and_then(|existing| existing.request_metadata.clone()) + } }); let request_body_ref = persisted_usage_body_ref( usage.request_body_ref.as_deref(), usage.request_body.as_ref(), + usage.request_body_state, request_metadata.as_ref(), existing.as_ref(), UsageBodyField::RequestBody, @@ -2804,6 +2925,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { let provider_request_body_ref = persisted_usage_body_ref( usage.provider_request_body_ref.as_deref(), usage.provider_request_body.as_ref(), + usage.provider_request_body_state, request_metadata.as_ref(), existing.as_ref(), UsageBodyField::ProviderRequestBody, @@ -2811,6 +2933,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { let response_body_ref = persisted_usage_body_ref( usage.response_body_ref.as_deref(), usage.response_body.as_ref(), + usage.response_body_state, request_metadata.as_ref(), existing.as_ref(), UsageBodyField::ResponseBody, @@ -2818,10 +2941,34 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { let client_response_body_ref = persisted_usage_body_ref( usage.client_response_body_ref.as_deref(), usage.client_response_body.as_ref(), + usage.client_response_body_state, request_metadata.as_ref(), existing.as_ref(), UsageBodyField::ClientResponseBody, ); + if clear_request_body + || clear_provider_request_body + || clear_response_body + || clear_client_response_body + { + let mut detached_bodies = self.detached_bodies.write().expect("usage repository lock"); + for (clear, field) in [ + (clear_request_body, UsageBodyField::RequestBody), + ( + clear_provider_request_body, + UsageBodyField::ProviderRequestBody, + ), + (clear_response_body, UsageBodyField::ResponseBody), + ( + clear_client_response_body, + UsageBodyField::ClientResponseBody, + ), + ] { + if clear { + detached_bodies.remove(&usage_body_ref(&usage.request_id, field)); + } + } + } let stored = StoredRequestUsageAudit { id: existing .as_ref() @@ -2935,11 +3082,15 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { .as_ref() .and_then(|existing| existing.request_headers.clone()) }), - request_body: usage.request_body.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.request_body.clone()) - }), + request_body: if clear_request_body { + None + } else { + usage.request_body.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.request_body.clone()) + }) + }, request_body_ref, request_body_state: usage.request_body_state.or_else(|| { existing @@ -2951,11 +3102,15 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { .as_ref() .and_then(|existing| existing.provider_request_headers.clone()) }), - provider_request_body: usage.provider_request_body.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.provider_request_body.clone()) - }), + provider_request_body: if clear_provider_request_body { + None + } else { + usage.provider_request_body.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.provider_request_body.clone()) + }) + }, provider_request_body_ref, provider_request_body_state: usage.provider_request_body_state.or_else(|| { existing @@ -2967,11 +3122,15 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { .as_ref() .and_then(|existing| existing.response_headers.clone()) }), - response_body: usage.response_body.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.response_body.clone()) - }), + response_body: if clear_response_body { + None + } else { + usage.response_body.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.response_body.clone()) + }) + }, response_body_ref, response_body_state: usage.response_body_state.or_else(|| { existing @@ -2983,61 +3142,95 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { .as_ref() .and_then(|existing| existing.client_response_headers.clone()) }), - client_response_body: usage.client_response_body.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.client_response_body.clone()) - }), + client_response_body: if clear_client_response_body { + None + } else { + usage.client_response_body.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.client_response_body.clone()) + }) + }, client_response_body_ref, client_response_body_state: usage.client_response_body_state.or_else(|| { existing .as_ref() .and_then(|existing| existing.client_response_body_state) }), - candidate_id: usage.candidate_id.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned)) - }), - candidate_index: usage.candidate_index.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_candidate_index()) - }), - key_name: usage.key_name.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_key_name().map(ToOwned::to_owned)) - }), - planner_kind: usage.planner_kind.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_planner_kind().map(ToOwned::to_owned)) - }), - route_family: usage.route_family.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_route_family().map(ToOwned::to_owned)) - }), - route_kind: usage.route_kind.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_route_kind().map(ToOwned::to_owned)) - }), - execution_path: usage.execution_path.or_else(|| { - existing - .as_ref() - .and_then(|existing| existing.routing_execution_path().map(ToOwned::to_owned)) - }), - local_execution_runtime_miss_reason: usage.local_execution_runtime_miss_reason.or_else( - || { + candidate_id: if replace_routing_snapshot { + usage.candidate_id + } else { + usage.candidate_id.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned)) + }) + }, + candidate_index: if replace_routing_snapshot { + usage.candidate_index + } else { + usage.candidate_index.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_candidate_index()) + }) + }, + key_name: if replace_routing_snapshot { + usage.key_name + } else { + usage.key_name.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_key_name().map(ToOwned::to_owned)) + }) + }, + planner_kind: if replace_routing_snapshot { + usage.planner_kind + } else { + usage.planner_kind.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_planner_kind().map(ToOwned::to_owned)) + }) + }, + route_family: if replace_routing_snapshot { + usage.route_family + } else { + usage.route_family.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_route_family().map(ToOwned::to_owned)) + }) + }, + route_kind: if replace_routing_snapshot { + usage.route_kind + } else { + usage.route_kind.or_else(|| { + existing + .as_ref() + .and_then(|existing| existing.routing_route_kind().map(ToOwned::to_owned)) + }) + }, + execution_path: if replace_routing_snapshot { + usage.execution_path + } else { + usage.execution_path.or_else(|| { + existing.as_ref().and_then(|existing| { + existing.routing_execution_path().map(ToOwned::to_owned) + }) + }) + }, + local_execution_runtime_miss_reason: if replace_routing_snapshot { + usage.local_execution_runtime_miss_reason + } else { + usage.local_execution_runtime_miss_reason.or_else(|| { existing.as_ref().and_then(|existing| { existing .routing_local_execution_runtime_miss_reason() .map(ToOwned::to_owned) }) - }, - ), + }) + }, client_family: usage_request_metadata_client_family(request_metadata.as_ref()) .map(ToOwned::to_owned) .or_else(|| { diff --git a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs index 14952c01c..1c9f76e08 100644 --- a/crates/aether-data/runtime/src/repository/usage/memory/tests.rs +++ b/crates/aether-data/runtime/src/repository/usage/memory/tests.rs @@ -15,7 +15,7 @@ use crate::repository::usage::{ }; use aether_data_contracts::repository::usage::{ usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy, - UsageAuditAggregationQuery, UsageBodyField, UsageDashboardSummaryQuery, + UsageAuditAggregationQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery, UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery, UsageTimeSeriesGranularity, }; @@ -135,6 +135,127 @@ fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord { } } +#[tokio::test] +async fn upsert_uses_typed_provider_capture_as_the_fast_fact_snapshot() { + for (name, state, incoming_tier, expected_tier) in [ + ( + "disabled-clear", + UsageBodyCaptureState::Disabled, + None, + None, + ), + ( + "truncated-clear", + UsageBodyCaptureState::Truncated, + None, + None, + ), + ( + "unavailable-clear", + UsageBodyCaptureState::Unavailable, + None, + None, + ), + ( + "disabled-preserve", + UsageBodyCaptureState::Disabled, + Some("priority"), + Some("priority"), + ), + ( + "truncated-preserve", + UsageBodyCaptureState::Truncated, + Some("priority"), + Some("priority"), + ), + ( + "none-clears-residual", + UsageBodyCaptureState::None, + Some("priority"), + None, + ), + ] { + let request_id = format!("req-memory-fast-{name}"); + let repository = InMemoryUsageReadRepository::default(); + let mut pending = sample_upsert_usage_record(&request_id); + pending.provider_request_body = Some(json!({ + "model": "gpt-5", + "service_tier": "priority" + })); + pending.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + pending.request_metadata = Some(json!({"provider_service_tier": "priority"})); + pending.target_model = Some("candidate-a-target".to_string()); + pending.candidate_id = Some("candidate-a".to_string()); + pending.candidate_index = Some(1); + pending.key_name = Some("key-a".to_string()); + pending.planner_kind = Some("planner-a".to_string()); + pending.route_family = Some("route-family-a".to_string()); + pending.route_kind = Some("route-kind-a".to_string()); + pending.execution_path = Some("path-a".to_string()); + let pending = repository + .upsert(pending) + .await + .expect("pending usage should upsert"); + assert_eq!(pending.provider_service_tier().as_deref(), Some("priority")); + + let mut terminal = sample_upsert_usage_record(&request_id); + terminal.status = "completed".to_string(); + terminal.provider_request_body_state = Some(state); + terminal.target_model = None; + terminal.request_metadata = + incoming_tier.map(|tier| json!({"provider_service_tier": tier})); + terminal.updated_at_unix_secs += 1; + terminal.finalized_at_unix_secs = Some(terminal.updated_at_unix_secs); + if state == UsageBodyCaptureState::None { + terminal.provider_request_body = Some(json!({ + "model": "stale-model", + "service_tier": "priority" + })); + terminal.provider_request_body_ref = Some(usage_body_ref( + &request_id, + UsageBodyField::ProviderRequestBody, + )); + } + let terminal = repository + .upsert(terminal) + .await + .expect("terminal usage should upsert"); + assert_eq!( + terminal.provider_service_tier().as_deref(), + expected_tier, + "state={state:?} must use only incoming facts" + ); + assert_eq!(terminal.target_model, None); + assert_eq!(terminal.candidate_id, None); + assert_eq!(terminal.candidate_index, None); + assert_eq!(terminal.key_name, None); + assert_eq!(terminal.planner_kind, None); + assert_eq!(terminal.route_family, None); + assert_eq!(terminal.route_kind, None); + assert_eq!(terminal.execution_path, None); + if state == UsageBodyCaptureState::None { + assert_eq!(terminal.provider_request_body, None); + assert_eq!(terminal.provider_request_body_ref, None); + + let mut late = sample_upsert_usage_record(&request_id); + late.provider_request_body = Some(json!({ + "model": "late-model", + "service_tier": "priority" + })); + late.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + late.request_metadata = Some(json!({"provider_service_tier": "priority"})); + late.updated_at_unix_secs += 2; + let late = repository + .upsert(late) + .await + .expect("late pending usage should not regress terminal capture"); + assert_eq!(late.provider_service_tier(), None); + assert_eq!(late.provider_request_body, None); + assert_eq!(late.provider_request_body_ref, None); + } + } +} + #[tokio::test] async fn finds_usage_by_request_id() { let repository = InMemoryUsageReadRepository::seed(vec![ diff --git a/crates/aether-usage/runtime/src/record.rs b/crates/aether-usage/runtime/src/record.rs index bcd9bca16..697a80204 100644 --- a/crates/aether-usage/runtime/src/record.rs +++ b/crates/aether-usage/runtime/src/record.rs @@ -2,7 +2,10 @@ use aether_data_contracts::repository::usage::UpsertUsageRecord; use aether_data_contracts::DataLayerError; use crate::request_metadata::{ - attach_provider_request_body_metadata, sanitize_usage_request_metadata, + attach_client_request_body_metadata, attach_provider_request_body_metadata, + clear_client_request_body_metadata, clear_provider_request_body_metadata, + request_body_derived_facts_action, sanitize_usage_request_metadata, + RequestBodyDerivedFactsAction, }; use crate::{UsageEvent, UsageEventType}; @@ -38,15 +41,42 @@ pub fn build_upsert_usage_record_from_event( } }; let mut data = event.data.clone(); - data.request_metadata = attach_provider_request_body_metadata( - data.request_metadata, - data.endpoint_api_format - .as_deref() - .or(data.api_format.as_deref()), - data.target_model.as_deref().or(Some(data.model.as_str())), - Some(data.model.as_str()), + // Request-derived facts are captured before body capture policy is applied. Do not let a + // truncation/disabled placeholder clear those facts while converting the queued event into a + // database record. Inline (or ref-loaded) bodies remain authoritative and may clear stale + // metadata when the final upstream request no longer contains a value. + match request_body_derived_facts_action(data.request_body.as_ref(), data.request_body_state) { + RequestBodyDerivedFactsAction::Refresh => { + data.request_metadata = attach_client_request_body_metadata( + data.request_metadata, + data.request_body.as_ref(), + ); + } + RequestBodyDerivedFactsAction::Clear => { + data.request_metadata = clear_client_request_body_metadata(data.request_metadata); + } + RequestBodyDerivedFactsAction::Preserve => {} + } + match request_body_derived_facts_action( data.provider_request_body.as_ref(), - ); + data.provider_request_body_state, + ) { + RequestBodyDerivedFactsAction::Refresh => { + data.request_metadata = attach_provider_request_body_metadata( + data.request_metadata, + data.endpoint_api_format + .as_deref() + .or(data.api_format.as_deref()), + data.target_model.as_deref().or(Some(data.model.as_str())), + Some(data.model.as_str()), + data.provider_request_body.as_ref(), + ); + } + RequestBodyDerivedFactsAction::Clear => { + data.request_metadata = clear_provider_request_body_metadata(data.request_metadata); + } + RequestBodyDerivedFactsAction::Preserve => {} + } let now_unix_secs = event.timestamp_ms / 1_000; Ok(UpsertUsageRecord { @@ -165,6 +195,8 @@ fn empty_to_none(value: Option) -> Option { #[cfg(test)] mod tests { + use aether_data_contracts::repository::usage::UsageBodyCaptureState; + use crate::{UsageEvent, UsageEventData, UsageEventType}; use super::build_upsert_usage_record_from_event; @@ -186,6 +218,9 @@ mod tests { output_tokens: Some(20), total_tokens: Some(30), status_code: Some(200), + request_body: Some(serde_json::json!({ + "reasoning": { "effort": "xhigh" } + })), provider_request_body: Some(serde_json::json!({ "reasoning": { "effort": "max" }, "service_tier": "priority" @@ -202,6 +237,14 @@ mod tests { assert_eq!(record.status, "completed"); assert_eq!(record.billing_status, "pending"); assert_eq!(record.total_tokens, Some(30)); + assert_eq!( + record + .request_metadata + .as_ref() + .and_then(|value| value.get("requested_reasoning_effort")) + .and_then(serde_json::Value::as_str), + Some("xhigh") + ); assert_eq!( record .request_metadata @@ -229,6 +272,139 @@ mod tests { assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000)); } + #[test] + fn truncated_request_placeholders_preserve_pre_capture_request_facts() { + let truncated = serde_json::json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded", + "max_bytes": 128, + "source_bytes": 2048, + "value_kind": "object" + }); + let record = build_upsert_usage_record_from_event(&UsageEvent { + event_type: UsageEventType::Completed, + request_id: "req-truncated-request-facts".to_string(), + timestamp_ms: 1_700_000_000_000, + data: UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + api_format: Some("openai:responses".to_string()), + endpoint_api_format: Some("openai:responses".to_string()), + request_body: Some(truncated.clone()), + request_body_state: Some(UsageBodyCaptureState::Truncated), + provider_request_body: Some(truncated), + provider_request_body_state: Some(UsageBodyCaptureState::Truncated), + request_metadata: Some(serde_json::json!({ + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max", + "provider_service_tier": "priority" + })), + ..UsageEventData::default() + }, + }) + .expect("record should build"); + + let metadata = record + .request_metadata + .as_ref() + .expect("derived request facts should remain"); + assert_eq!(metadata["requested_reasoning_effort"], "xhigh"); + assert_eq!(metadata["provider_reasoning_effort"], "max"); + assert_eq!(metadata["provider_service_tier"], "priority"); + } + + #[test] + fn disabled_reference_and_unavailable_bodies_preserve_pre_capture_request_facts() { + for state in [ + UsageBodyCaptureState::Disabled, + UsageBodyCaptureState::Reference, + UsageBodyCaptureState::Unavailable, + ] { + let record = build_upsert_usage_record_from_event(&UsageEvent { + event_type: UsageEventType::Completed, + request_id: format!("req-preserved-{state:?}"), + timestamp_ms: 1_700_000_000_000, + data: UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + request_body_state: Some(state), + provider_request_body_state: Some(state), + request_metadata: Some(serde_json::json!({ + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max", + "provider_service_tier": "priority" + })), + ..UsageEventData::default() + }, + }) + .expect("record should build"); + + let metadata = record + .request_metadata + .as_ref() + .expect("derived request facts should remain"); + assert_eq!(metadata["requested_reasoning_effort"], "xhigh"); + assert_eq!(metadata["provider_reasoning_effort"], "max"); + assert_eq!(metadata["provider_service_tier"], "priority"); + } + } + + #[test] + fn missing_or_complete_factless_final_bodies_clear_stale_request_facts() { + let cases = [ + ("typed-missing", Some(UsageBodyCaptureState::None), None), + ( + "typed-missing-with-stale-body", + Some(UsageBodyCaptureState::None), + Some(serde_json::json!({ + "reasoning_effort": "xhigh", + "service_tier": "priority" + })), + ), + ("untyped-missing", None, None), + ( + "inline-without-facts", + Some(UsageBodyCaptureState::Inline), + Some(serde_json::json!({"model": "gpt-5"})), + ), + ]; + + for (name, state, body) in cases { + let record = build_upsert_usage_record_from_event(&UsageEvent { + event_type: UsageEventType::Completed, + request_id: format!("req-clear-{name}"), + timestamp_ms: 1_700_000_000_000, + data: UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + request_body: body.clone(), + request_body_state: state, + provider_request_body: body, + provider_request_body_state: state, + request_metadata: Some(serde_json::json!({ + "trace_id": "trace-1", + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max", + "provider_service_tier": "priority", + "provider_actual_service_tier": "priority" + })), + ..UsageEventData::default() + }, + }) + .expect("record should build"); + + let metadata = record + .request_metadata + .as_ref() + .expect("audit facts remain"); + assert_eq!(metadata["trace_id"], "trace-1"); + assert_eq!(metadata["provider_actual_service_tier"], "priority"); + assert!(metadata.get("requested_reasoning_effort").is_none()); + assert!(metadata.get("provider_reasoning_effort").is_none()); + assert!(metadata.get("provider_service_tier").is_none()); + } + } + #[test] fn cancelled_terminal_record_is_void_for_billing() { let record = build_upsert_usage_record_from_event(&UsageEvent { diff --git a/crates/aether-usage/runtime/src/request_metadata.rs b/crates/aether-usage/runtime/src/request_metadata.rs index 4c752e4b1..5d830acf7 100644 --- a/crates/aether-usage/runtime/src/request_metadata.rs +++ b/crates/aether-usage/runtime/src/request_metadata.rs @@ -6,9 +6,10 @@ use aether_contracts::ExecutionPlan; use aether_data_contracts::repository::usage::{ extract_provider_actual_service_tier_from_response, extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body, - normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, + normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, UsageBodyCaptureState, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, + REQUESTED_REASONING_EFFORT_METADATA_KEY, }; use serde_json::{json, Map, Value}; @@ -17,6 +18,53 @@ const MAX_USAGE_REQUEST_METADATA_NODES: usize = 4_000; const MAX_USAGE_REQUEST_METADATA_BYTES: usize = 16 * 1024; const MAX_USAGE_REQUEST_METADATA_STRING_BYTES: usize = 1_024; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RequestBodyDerivedFactsAction { + Refresh, + Preserve, + Clear, +} + +pub(crate) fn request_body_derived_facts_action( + request_body: Option<&Value>, + state: Option, +) -> RequestBodyDerivedFactsAction { + // A typed `none` marker is produced by the final capture attempt. It must take precedence over + // a body value that may have survived from an earlier candidate. A missing marker (`None`) + // remains compatible with legacy events, where a present body is still authoritative. + if state == Some(UsageBodyCaptureState::None) { + return RequestBodyDerivedFactsAction::Clear; + } + + if let Some(request_body) = request_body { + if matches!( + state, + Some( + UsageBodyCaptureState::Truncated + | UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable + ) + ) || 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") + }) { + return RequestBodyDerivedFactsAction::Preserve; + } + return RequestBodyDerivedFactsAction::Refresh; + } + + match state { + Some( + UsageBodyCaptureState::Inline + | UsageBodyCaptureState::Reference + | UsageBodyCaptureState::Truncated + | UsageBodyCaptureState::Disabled + | UsageBodyCaptureState::Unavailable, + ) => RequestBodyDerivedFactsAction::Preserve, + Some(UsageBodyCaptureState::None) | None => RequestBodyDerivedFactsAction::Clear, + } +} + pub(crate) fn build_usage_request_metadata_seed( _plan: &ExecutionPlan, context: Option<&Map>, @@ -76,6 +124,32 @@ pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Opti (!filtered.is_empty()).then_some(Value::Object(filtered)) } +pub(crate) fn attach_client_request_body_metadata( + metadata: Option, + request_body: Option<&Value>, +) -> Option { + let request_body_is_object = request_body.and_then(Value::as_object).is_some(); + let reasoning_effort = extract_provider_reasoning_effort_from_body(request_body); + if !request_body_is_object && reasoning_effort.is_none() { + return metadata; + } + + let mut object = match metadata { + Some(Value::Object(object)) => object, + _ => Map::new(), + }; + if request_body_is_object { + object.remove(REQUESTED_REASONING_EFFORT_METADATA_KEY); + } + if let Some(reasoning_effort) = reasoning_effort { + object.insert( + REQUESTED_REASONING_EFFORT_METADATA_KEY.to_string(), + Value::String(reasoning_effort), + ); + } + (!object.is_empty()).then_some(Value::Object(object)) +} + pub(crate) fn attach_provider_request_body_metadata( metadata: Option, provider_api_format: Option<&str>, @@ -129,6 +203,31 @@ pub(crate) fn attach_provider_request_body_metadata( (!object.is_empty()).then_some(Value::Object(object)) } +pub(crate) fn clear_client_request_body_metadata(metadata: Option) -> Option { + clear_request_metadata_fields(metadata, &[REQUESTED_REASONING_EFFORT_METADATA_KEY]) +} + +pub(crate) fn clear_provider_request_body_metadata(metadata: Option) -> Option { + clear_request_metadata_fields( + metadata, + &[ + PROVIDER_REASONING_EFFORT_METADATA_KEY, + PROVIDER_SERVICE_TIER_METADATA_KEY, + PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, + ], + ) +} + +fn clear_request_metadata_fields(metadata: Option, keys: &[&str]) -> Option { + let Some(Value::Object(mut object)) = metadata else { + return None; + }; + for key in keys { + object.remove(*key); + } + (!object.is_empty()).then_some(Value::Object(object)) +} + pub(crate) fn attach_provider_response_body_metadata( metadata: Option, provider_response_body: Option<&Value>, @@ -148,6 +247,47 @@ pub(crate) fn attach_provider_response_body_metadata( attach_provider_actual_service_tier_metadata(metadata, actual_service_tier.as_deref()) } +/// Refreshes the response-derived tier for a terminal snapshot. Complete response objects are +/// authoritative even when they contain no tier (which clears a stale candidate value). Capture +/// placeholders/absent bodies are not authoritative, so a terminal summary already present in +/// metadata is preserved for those cases. +pub(crate) fn refresh_provider_response_body_metadata( + metadata: Option, + provider_response_body: Option<&Value>, +) -> Option { + let is_capture_placeholder = provider_response_body + .and_then(Value::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") + }); + let body_is_complete_object = + provider_response_body.and_then(Value::as_object).is_some() && !is_capture_placeholder; + let actual_service_tier = + extract_provider_actual_service_tier_from_response(provider_response_body) + .and_then(|value| normalize_provider_service_tier(&value)); + let Some(actual_service_tier) = actual_service_tier else { + if !body_is_complete_object { + return metadata; + } + let Some(Value::Object(mut object)) = metadata else { + return None; + }; + object.remove(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY); + return (!object.is_empty()).then_some(Value::Object(object)); + }; + + let mut object = match metadata { + Some(Value::Object(object)) => object, + _ => Map::new(), + }; + object.insert( + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY.to_string(), + Value::String(actual_service_tier), + ); + (!object.is_empty()).then_some(Value::Object(object)) +} + pub(crate) fn attach_provider_actual_service_tier_metadata( metadata: Option, actual_service_tier: Option<&str>, @@ -179,6 +319,7 @@ fn copy_allowed_metadata_fields(source: &Map, target: &mut Map, target: &mut Map remove_non_empty_string(&mut source, target, "request_path"); remove_non_empty_string(&mut source, target, "request_query_string"); remove_non_empty_string(&mut source, target, "request_path_and_query"); + remove_non_empty_string(&mut source, target, REQUESTED_REASONING_EFFORT_METADATA_KEY); remove_non_empty_string(&mut source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY); remove_non_empty_string(&mut source, target, PROVIDER_SERVICE_TIER_METADATA_KEY); remove_non_empty_string( @@ -517,10 +659,16 @@ mod tests { use serde_json::{json, Value}; use std::collections::BTreeMap; + use crate::{ + apply_usage_body_capture_policy_to_event, UsageBodyCapturePolicy, UsageEvent, + UsageEventData, UsageEventType, UsageRequestRecordLevel, + }; + use super::{ - attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata, - attach_provider_response_body_metadata, build_usage_request_metadata_seed, - merge_usage_request_metadata, merge_usage_request_metadata_owned, + attach_client_request_body_metadata, attach_provider_actual_service_tier_metadata, + attach_provider_request_body_metadata, attach_provider_response_body_metadata, + build_usage_request_metadata_seed, merge_usage_request_metadata, + merge_usage_request_metadata_owned, refresh_provider_response_body_metadata, sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES, @@ -830,6 +978,27 @@ mod tests { assert_eq!(metadata, None); } + #[test] + fn client_request_body_metadata_preserves_requested_reasoning_mapping_source() { + let updated = attach_client_request_body_metadata( + Some(json!({ + "trace_id": "trace-1", + "requested_reasoning_effort": "low" + })), + Some(&json!({ + "reasoning": { "effort": "XHigh" } + })), + ) + .expect("metadata should remain"); + + assert_eq!(updated["requested_reasoning_effort"], "xhigh"); + + let cleared = + attach_client_request_body_metadata(Some(updated), Some(&json!({ "model": "gpt-5" }))) + .expect("trace metadata should remain"); + assert!(cleared.get("requested_reasoning_effort").is_none()); + } + #[test] fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() { let metadata = Some(json!({ @@ -880,6 +1049,75 @@ mod tests { ); } + #[test] + fn final_provider_request_tier_survives_basic_body_capture_as_derived_metadata() { + let request_body = json!({ + "model": "gpt-5", + "reasoning": { "effort": "xhigh" } + }); + let provider_request_body = json!({ + "model": "gpt-5", + "reasoning": { "effort": "max" }, + "service_tier": "priority" + }); + let request_metadata = attach_client_request_body_metadata(None, Some(&request_body)); + let request_metadata = attach_provider_request_body_metadata( + request_metadata, + Some("openai:responses"), + Some("gpt-5"), + Some("gpt-5"), + Some(&provider_request_body), + ); + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-final-provider-tier", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + request_body: Some(request_body), + provider_request_body: Some(provider_request_body), + request_metadata, + ..UsageEventData::default() + }, + ); + + apply_usage_body_capture_policy_to_event( + UsageBodyCapturePolicy { + record_level: UsageRequestRecordLevel::Basic, + max_request_body_bytes: Some(1024), + max_response_body_bytes: Some(1024), + }, + &mut event, + ); + + assert_eq!(event.data.request_body, None); + assert_eq!(event.data.provider_request_body, None); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("requested_reasoning_effort")), + Some(&json!("xhigh")) + ); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_reasoning_effort")), + Some(&json!("max")) + ); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_service_tier")), + Some(&json!("priority")) + ); + } + #[test] fn provider_response_metadata_preserves_terminal_actual_service_tier() { let metadata = attach_provider_response_body_metadata( @@ -902,6 +1140,39 @@ mod tests { ); } + #[test] + fn terminal_response_refresh_replaces_stale_actual_tier() { + let metadata = refresh_provider_response_body_metadata( + Some(json!({"provider_actual_service_tier": "priority"})), + Some(&json!({"service_tier": "Default"})), + ) + .expect("actual tier should remain"); + assert_eq!(metadata["provider_actual_service_tier"], "default"); + + let metadata = refresh_provider_response_body_metadata( + Some(json!({ + "trace_id": "trace-1", + "provider_actual_service_tier": "priority" + })), + Some(&json!({"id": "response-without-tier"})), + ) + .expect("un-tiered complete response should remain auditable"); + assert_eq!(metadata, json!({"trace_id": "trace-1"})); + } + + #[test] + fn terminal_response_refresh_preserves_summary_over_capture_placeholder() { + let metadata = refresh_provider_response_body_metadata( + Some(json!({"provider_actual_service_tier": "priority"})), + Some(&json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded" + })), + ) + .expect("summary should survive placeholder capture"); + assert_eq!(metadata["provider_actual_service_tier"], "priority"); + } + #[test] fn terminal_summary_tier_uses_the_same_normalized_metadata_field() { let metadata = attach_provider_actual_service_tier_metadata( diff --git a/crates/aether-usage/runtime/src/runtime.rs b/crates/aether-usage/runtime/src/runtime.rs index ddf350382..542270b4d 100644 --- a/crates/aether-usage/runtime/src/runtime.rs +++ b/crates/aether-usage/runtime/src/runtime.rs @@ -13,7 +13,12 @@ use tokio::sync::mpsc; use tracing::{info, warn}; use crate::executor::spawn_on_usage_background_runtime; -use crate::request_metadata::attach_provider_response_body_metadata; +use crate::request_metadata::{ + attach_client_request_body_metadata, attach_provider_request_body_metadata, + attach_provider_response_body_metadata, clear_client_request_body_metadata, + clear_provider_request_body_metadata, request_body_derived_facts_action, + RequestBodyDerivedFactsAction, +}; use crate::worker::{ build_usage_queue_worker_with_record_gate, UsageWorkerControl, UsageWorkerObservation, }; @@ -1264,6 +1269,7 @@ impl UsageRuntime { where T: UsageRuntimeAccess, { + preserve_request_facts(event); preserve_provider_response_facts(event); match self.cached_body_capture_policy(data).await { Ok(policy) => apply_usage_body_capture_policy_to_event(policy, event), @@ -1734,6 +1740,44 @@ impl UsageRuntime { } } +fn preserve_request_facts(event: &mut UsageEvent) { + let data = &mut event.data; + match request_body_derived_facts_action(data.request_body.as_ref(), data.request_body_state) { + RequestBodyDerivedFactsAction::Refresh => { + data.request_metadata = attach_client_request_body_metadata( + data.request_metadata.take(), + data.request_body.as_ref(), + ); + } + RequestBodyDerivedFactsAction::Clear => { + data.request_metadata = + clear_client_request_body_metadata(data.request_metadata.take()); + } + RequestBodyDerivedFactsAction::Preserve => {} + } + match request_body_derived_facts_action( + data.provider_request_body.as_ref(), + data.provider_request_body_state, + ) { + RequestBodyDerivedFactsAction::Refresh => { + data.request_metadata = attach_provider_request_body_metadata( + data.request_metadata.take(), + data.endpoint_api_format + .as_deref() + .or(data.api_format.as_deref()), + data.target_model.as_deref().or(Some(data.model.as_str())), + Some(data.model.as_str()), + data.provider_request_body.as_ref(), + ); + } + RequestBodyDerivedFactsAction::Clear => { + data.request_metadata = + clear_provider_request_body_metadata(data.request_metadata.take()); + } + RequestBodyDerivedFactsAction::Preserve => {} + } +} + fn preserve_provider_response_facts(event: &mut UsageEvent) { let metadata = event.data.request_metadata.take(); event.data.request_metadata = @@ -2989,7 +3033,9 @@ mod tests { use aether_data_contracts::repository::settlement::{ StoredUsageSettlement, UsageSettlementInput, }; - use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord}; + use aether_data_contracts::repository::usage::{ + StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, + }; use aether_data_contracts::DataLayerError; use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState}; use async_trait::async_trait; @@ -2998,9 +3044,10 @@ mod tests { use tokio::time::{sleep, timeout, Duration}; use super::{ - preserve_provider_response_facts, LifecycleEventCoalescer, UsageBillingEventEnricher, - UsageBodyCapturePolicy, UsageEnqueueRetryDispatcher, UsageRequestRecordLevel, - UsageRuntimeAccess, UsageWorkerObservation, UsageWorkerSupervisorState, + preserve_provider_response_facts, preserve_request_facts, LifecycleEventCoalescer, + UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEnqueueRetryDispatcher, + UsageRequestRecordLevel, UsageRuntimeAccess, UsageWorkerObservation, + UsageWorkerSupervisorState, }; use crate::worker::ManualProxyNodeCounter; use crate::{ @@ -5729,6 +5776,76 @@ mod tests { ); } + #[test] + fn preserve_request_facts_ignores_post_capture_truncation_placeholders() { + let truncated = json!({ + "truncated": true, + "reason": "body_capture_limit_exceeded" + }); + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-truncated-preserve", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + request_body: Some(truncated.clone()), + request_body_state: Some(UsageBodyCaptureState::Truncated), + provider_request_body: Some(truncated), + provider_request_body_state: Some(UsageBodyCaptureState::Truncated), + request_metadata: Some(json!({ + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max", + "provider_service_tier": "priority" + })), + ..UsageEventData::default() + }, + ); + + preserve_request_facts(&mut event); + + let metadata = event + .data + .request_metadata + .as_ref() + .expect("metadata remains"); + assert_eq!(metadata["requested_reasoning_effort"], "xhigh"); + assert_eq!(metadata["provider_reasoning_effort"], "max"); + assert_eq!(metadata["provider_service_tier"], "priority"); + } + + #[test] + fn preserve_request_facts_clears_stale_facts_when_final_bodies_are_missing() { + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-missing-final-body", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5".to_string(), + request_body_state: Some(UsageBodyCaptureState::None), + provider_request_body_state: Some(UsageBodyCaptureState::None), + request_metadata: Some(json!({ + "requested_reasoning_effort": "xhigh", + "provider_reasoning_effort": "max", + "provider_service_tier": "priority", + "provider_actual_service_tier": "priority" + })), + ..UsageEventData::default() + }, + ); + + preserve_request_facts(&mut event); + + let metadata = event + .data + .request_metadata + .as_ref() + .expect("response audit fact remains"); + assert!(metadata.get("requested_reasoning_effort").is_none()); + assert!(metadata.get("provider_reasoning_effort").is_none()); + assert!(metadata.get("provider_service_tier").is_none()); + assert_eq!(metadata["provider_actual_service_tier"], "priority"); + } + #[test] fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() { let mut event = UsageEvent::new( @@ -5739,9 +5856,16 @@ mod tests { model: "gpt-5".to_string(), total_tokens: Some(42), error_message: Some("upstream failed".to_string()), - request_body: Some(json!({"messages":[{"role":"user","content":"hello"}]})), + request_body: Some(json!({ + "messages":[{"role":"user","content":"hello"}], + "reasoning": {"effort": "xhigh"} + })), request_body_ref: Some("usage://request/req-basic-1/request_body".to_string()), - provider_request_body: Some(json!({"model":"gpt-5"})), + provider_request_body: Some(json!({ + "model":"gpt-5", + "reasoning": {"effort": "max"}, + "service_tier": "priority" + })), provider_request_body_ref: Some( "usage://request/req-basic-1/provider_request_body".to_string(), ), @@ -5754,11 +5878,16 @@ mod tests { client_response_body_ref: Some( "usage://request/req-basic-1/client_response_body".to_string(), ), - request_metadata: Some(json!({"provider_service_tier": "priority"})), + request_metadata: Some(json!({ + "requested_reasoning_effort": "low", + "provider_reasoning_effort": "medium", + "provider_service_tier": "standard" + })), ..UsageEventData::default() }, ); + preserve_request_facts(&mut event); preserve_provider_response_facts(&mut event); apply_usage_body_capture_policy_to_event( UsageBodyCapturePolicy { @@ -5778,6 +5907,33 @@ mod tests { assert!(event.data.response_body_ref.is_none()); assert!(event.data.client_response_body.is_none()); assert!(event.data.client_response_body_ref.is_none()); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("requested_reasoning_effort")) + .and_then(serde_json::Value::as_str), + Some("xhigh") + ); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_reasoning_effort")) + .and_then(serde_json::Value::as_str), + Some("max") + ); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_service_tier")) + .and_then(serde_json::Value::as_str), + Some("priority") + ); assert_eq!( event .data diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index f2ab5ca93..bb644d8ee 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -13,10 +13,11 @@ use crate::body_capture::{ RuntimeBodyCaptureMetadataInput, }; use crate::request_metadata::{ - attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata, - attach_provider_response_body_metadata, build_usage_request_metadata_seed, + attach_client_request_body_metadata, attach_provider_actual_service_tier_metadata, + attach_provider_request_body_metadata, build_usage_request_metadata_seed, merge_usage_request_metadata, merge_usage_request_metadata_owned, - sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, + refresh_provider_response_body_metadata, sanitize_usage_request_metadata, + sanitize_usage_request_metadata_ref, }; use crate::{ map_usage_from_response, stream_capture_terminal_state, GatewayStreamReportRequest, @@ -696,6 +697,8 @@ fn build_terminal_usage_event_from_seed_impl( } else { merge_usage_request_metadata(request_metadata, audit_payload) }; + let request_metadata = + attach_client_request_body_metadata(request_metadata, request_body.as_ref()); let request_metadata = attach_provider_request_body_metadata( request_metadata, Some(provider_contract.as_str()), @@ -1021,7 +1024,7 @@ pub fn build_sync_terminal_usage_seed( status_code, provider_response_full.as_ref(), ); - let request_metadata = attach_provider_response_body_metadata( + let request_metadata = refresh_provider_response_body_metadata( context_seed.request_metadata, provider_response_full.as_ref(), ); @@ -1199,12 +1202,16 @@ pub fn build_stream_terminal_usage_seed( missing_observed_finish, terminal_error_message.is_some(), ); - let request_metadata = attach_provider_actual_service_tier_metadata( + let request_metadata = refresh_provider_response_body_metadata( context_seed.request_metadata, + provider_response_full.as_ref(), + ); + // The parser's terminal summary is authoritative when a response body is truncated or the + // body and summary disagree; attach it after the body refresh so it wins. + let request_metadata = attach_provider_actual_service_tier_metadata( + request_metadata, provider_actual_service_tier.as_deref(), ); - let request_metadata = - attach_provider_response_body_metadata(request_metadata, provider_response_full.as_ref()); TerminalUsageSeed { terminal_state, @@ -2111,7 +2118,10 @@ fn build_runtime_request_metadata_seed_from_parts( metadata.insert("body_size".to_string(), body_size); } - (!metadata.is_empty()).then_some(Value::Object(metadata)) + attach_client_request_body_metadata( + (!metadata.is_empty()).then_some(Value::Object(metadata)), + context_value_ref(context, "original_request_body"), + ) } fn build_runtime_body_size_metadata( diff --git a/frontend/src/api/__tests__/provider-keys.spec.ts b/frontend/src/api/__tests__/provider-keys.spec.ts new file mode 100644 index 000000000..5fecd1f93 --- /dev/null +++ b/frontend/src/api/__tests__/provider-keys.spec.ts @@ -0,0 +1,38 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { getMock } = vi.hoisted(() => ({ getMock: vi.fn() })) + +vi.mock('@/api/client', () => ({ + default: { + get: getMock, + }, +})) + +import { getProviderKeysPage } from '@/api/endpoints/keys' + +describe('getProviderKeysPage', () => { + beforeEach(() => { + getMock.mockReset() + }) + + it('normalizes a legacy array response for the provider drawer', async () => { + getMock.mockResolvedValue({ + data: [{ id: 'key-1' }, { id: 'key-2' }], + }) + + const result = await getProviderKeysPage('provider-demo', { page: 1, page_size: 1 }) + + expect(result).toMatchObject({ total: 2, page: 1, page_size: 1 }) + expect(result.keys).toEqual([{ id: 'key-1' }]) + }) + + it('normalizes a malformed object without exposing a non-array keys field', async () => { + getMock.mockResolvedValue({ + data: { total: null, page: null, page_size: null, keys: {} }, + }) + + const result = await getProviderKeysPage('provider-demo', { page: 2, page_size: 3 }) + + expect(result).toEqual({ total: 0, page: 2, page_size: 3, keys: [] }) + }) +}) diff --git a/frontend/src/api/__tests__/providers-summary.spec.ts b/frontend/src/api/__tests__/providers-summary.spec.ts new file mode 100644 index 000000000..9d7841605 --- /dev/null +++ b/frontend/src/api/__tests__/providers-summary.spec.ts @@ -0,0 +1,101 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { getMock } = vi.hoisted(() => ({ getMock: vi.fn() })) + +vi.mock('@/api/client', () => ({ + default: { + get: getMock, + }, +})) + +import { getProviderMappingPreview, getProvidersSummary } from '@/api/endpoints/providers' + +const provider = { + id: 'provider-1', + name: 'Provider 1', + provider_type: 'openai', + is_active: true, + endpoints: [], +} + +describe('getProvidersSummary', () => { + beforeEach(() => { + getMock.mockReset() + }) + + it('normalizes the paginated summary response', async () => { + getMock.mockResolvedValue({ + data: { total: 1, page: 1, page_size: 20, items: [provider] }, + }) + + const result = await getProvidersSummary({ page: 1, page_size: 20, search: 'paged' }) + + expect(result.total).toBe(1) + expect(result.items).toHaveLength(1) + expect(result.items[0]?.kiro_simulated_cache_enabled).toBe(false) + }) + + it('supports the legacy array response without reading an undefined items field', async () => { + getMock.mockResolvedValue({ data: [provider] }) + + const result = await getProvidersSummary({ page: 2, page_size: 20, search: 'legacy' }) + + expect(result).toMatchObject({ total: 1, page: 2, page_size: 20 }) + expect(result.items).toHaveLength(1) + }) +}) + +describe('getProviderMappingPreview', () => { + beforeEach(() => { + getMock.mockReset() + }) + + it('normalizes a non-contract response instead of exposing missing arrays to the UI', async () => { + getMock.mockResolvedValue({ + data: { message: '演示模式：该接口暂未模拟', demo_mode: true }, + }) + + const result = await getProviderMappingPreview('provider-demo') + + expect(result).toEqual({ + provider_id: 'provider-demo', + provider_name: '', + keys: [], + total_keys: 0, + total_matches: 0, + truncated: false, + truncated_keys: 0, + truncated_models: 0, + }) + }) + + it('normalizes missing nested mapping arrays', async () => { + getMock.mockResolvedValue({ + data: { + provider_id: 'provider-nested', + provider_name: 'Nested Provider', + keys: [{ + key_id: 'key-1', + key_name: 'Primary', + masked_key: 'sk-***', + is_active: true, + allowed_models: null, + matching_global_models: [{ + global_model_id: 'model-1', + global_model_name: 'gpt-5', + display_name: 'GPT-5', + is_active: true, + matched_models: null, + }], + }], + }, + }) + + const result = await getProviderMappingPreview('provider-nested') + + expect(result.keys[0]?.allowed_models).toEqual([]) + expect(result.keys[0]?.matching_global_models[0]?.matched_models).toEqual([]) + expect(result.total_keys).toBe(1) + expect(result.total_matches).toBe(1) + }) +}) diff --git a/frontend/src/api/__tests__/users.spec.ts b/frontend/src/api/__tests__/users.spec.ts index 52d29f3ae..3f2f8d486 100644 --- a/frontend/src/api/__tests__/users.spec.ts +++ b/frontend/src/api/__tests__/users.spec.ts @@ -1,13 +1,15 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { getMock, cachedRequestMock } = vi.hoisted(() => ({ +const { getMock, postMock, cachedRequestMock } = vi.hoisted(() => ({ getMock: vi.fn(), + postMock: vi.fn(), cachedRequestMock: vi.fn(async (_key: string, fn: () => Promise) => fn()), })) vi.mock('@/api/client', () => ({ default: { get: getMock, + post: postMock, }, })) @@ -20,6 +22,7 @@ import { usersApi } from '@/api/users' describe('usersApi admin list query', () => { beforeEach(() => { getMock.mockReset() + postMock.mockReset() cachedRequestMock.mockClear() getMock.mockResolvedValue({ data: { @@ -49,4 +52,55 @@ describe('usersApi admin list query', () => { }, }) }) + + it('keeps user management renderable when the group response has no items array', async () => { + getMock.mockResolvedValueOnce({ + data: { + message: '演示模式：该接口暂未模拟', + demo_mode: true, + }, + }) + + await expect(usersApi.listUserGroups()).resolves.toEqual({ + message: '演示模式：该接口暂未模拟', + demo_mode: true, + items: [], + }) + }) + + it('creates a managed key through the selected target user route', async () => { + postMock.mockResolvedValueOnce({ + data: { + id: 'target-key', + key: 'sk-target', + }, + }) + const payload = { + name: 'target key', + feature_settings: { + chat_pii_redaction: { enabled: true }, + }, + } + + await usersApi.createApiKey('target-user', payload) + + expect(postMock).toHaveBeenCalledWith( + '/api/admin/users/target-user/api-keys', + payload, + ) + }) + + it('reads managed keys from the production api_keys envelope', async () => { + getMock.mockResolvedValueOnce({ + data: { + api_keys: [{ id: 'target-key', created_at: '2026-07-17T00:00:00Z' }], + total: 1, + }, + }) + + await expect(usersApi.getUserApiKeys('target-user')).resolves.toEqual([ + { id: 'target-key', created_at: '2026-07-17T00:00:00Z' }, + ]) + expect(getMock).toHaveBeenCalledWith('/api/admin/users/target-user/api-keys') + }) }) diff --git a/frontend/src/api/dashboard.ts b/frontend/src/api/dashboard.ts index ab006747f..2f0e28372 100644 --- a/frontend/src/api/dashboard.ts +++ b/frontend/src/api/dashboard.ts @@ -209,6 +209,7 @@ export interface RequestDetail { has_format_conversion?: boolean | null model: string target_model?: string | null // 映射后的目标模型名 + requested_reasoning_effort?: string | null reasoning_effort?: string | null service_tier?: string | null actual_service_tier?: string | null @@ -266,6 +267,7 @@ export interface RequestDetail { response_time_ms: number first_byte_time_ms?: number | null created_at: string + updated_at?: string | null request_headers?: Record request_body?: Record provider_request_headers?: Record diff --git a/frontend/src/api/endpoints/keys.ts b/frontend/src/api/endpoints/keys.ts index aa4e31eba..08006df60 100644 --- a/frontend/src/api/endpoints/keys.ts +++ b/frontend/src/api/endpoints/keys.ts @@ -121,17 +121,50 @@ export interface ProviderKeysPageQuery { page_size?: number } +type ProviderKeysPagePayload = ProviderKeysPageResponse | EndpointAPIKey[] + +function normalizeProviderKeysPage( + value: ProviderKeysPagePayload, + page: number, + pageSize: number, +): ProviderKeysPageResponse { + if (Array.isArray(value)) { + const start = value.length > pageSize ? (page - 1) * pageSize : 0 + const keys = value.slice(start, start + pageSize) + return { + total: value.length, + page, + page_size: pageSize, + keys, + } + } + + const keys = Array.isArray(value.keys) ? value.keys : [] + return { + total: typeof value.total === 'number' && Number.isFinite(value.total) + ? value.total + : keys.length, + page: typeof value.page === 'number' && Number.isFinite(value.page) + ? value.page + : page, + page_size: typeof value.page_size === 'number' && Number.isFinite(value.page_size) + ? value.page_size + : pageSize, + keys, + } +} + export async function getProviderKeysPage( providerId: string, params: ProviderKeysPageQuery = {}, ): Promise { const page = params.page ?? 1 const pageSize = params.page_size ?? 20 - const response = await client.get( + const response = await client.get( `/api/admin/endpoints/providers/${providerId}/keys`, { params: { page, page_size: pageSize } }, ) - return response.data + return normalizeProviderKeysPage(response.data, page, pageSize) } export async function getProviderKeys(providerId: string): Promise { diff --git a/frontend/src/api/endpoints/providers.ts b/frontend/src/api/endpoints/providers.ts index 9e13c1d3a..c6a9f6dd5 100644 --- a/frontend/src/api/endpoints/providers.ts +++ b/frontend/src/api/endpoints/providers.ts @@ -42,6 +42,8 @@ export interface ProviderSummaryPageResponse { items: ProviderWithEndpointsSummary[] } +type ProviderSummaryResponse = ProviderSummaryPageResponse | ProviderWithEndpointsSummary[] + function normalizeProviderSummary( provider: ProviderWithEndpointsSummary, ): ProviderWithEndpointsSummary { @@ -62,16 +64,26 @@ export async function getProvidersSummary( return cachedRequest( cacheKey, async () => { - const response = await client.get( + const response = await client.get( '/api/admin/providers/summary', { params, timeout: options.timeout, }, ) + const data = response.data + if (Array.isArray(data)) { + return { + total: data.length, + page: params.page ?? 1, + page_size: params.page_size ?? data.length, + items: data.map(normalizeProviderSummary), + } + } + return { - ...response.data, - items: response.data.items.map(normalizeProviderSummary), + ...data, + items: (data.items ?? []).map(normalizeProviderSummary), } }, cacheTtlMs, @@ -371,6 +383,84 @@ export interface ProviderMappingPreviewResponse { truncated_models: number } +function mappingPreviewRecord(value: unknown): Record { + return value !== null && typeof value === 'object' && !Array.isArray(value) + ? value as Record + : {} +} + +function mappingPreviewString(value: unknown, fallback = ''): string { + return typeof value === 'string' ? value : fallback +} + +function mappingPreviewCount(value: unknown, fallback: number): number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 + ? value + : fallback +} + +function normalizeProviderMappingPreview( + value: unknown, + providerId: string, +): ProviderMappingPreviewResponse { + const source = mappingPreviewRecord(value) + const rawKeys = Array.isArray(source.keys) ? source.keys : [] + const keys = rawKeys.map((rawKey) => { + const key = mappingPreviewRecord(rawKey) + const rawGlobalModels = Array.isArray(key.matching_global_models) + ? key.matching_global_models + : [] + + return { + key_id: mappingPreviewString(key.key_id), + key_name: mappingPreviewString(key.key_name), + masked_key: mappingPreviewString(key.masked_key, '***'), + is_active: key.is_active === true, + allowed_models: Array.isArray(key.allowed_models) + ? key.allowed_models.filter((item): item is string => typeof item === 'string') + : [], + matching_global_models: rawGlobalModels.map((rawGlobalModel) => { + const globalModel = mappingPreviewRecord(rawGlobalModel) + const rawMatchedModels = Array.isArray(globalModel.matched_models) + ? globalModel.matched_models + : [] + + return { + global_model_id: mappingPreviewString(globalModel.global_model_id), + global_model_name: mappingPreviewString(globalModel.global_model_name), + display_name: mappingPreviewString( + globalModel.display_name, + mappingPreviewString(globalModel.global_model_name), + ), + is_active: globalModel.is_active === true, + matched_models: rawMatchedModels.map((rawMatchedModel) => { + const matchedModel = mappingPreviewRecord(rawMatchedModel) + return { + allowed_model: mappingPreviewString(matchedModel.allowed_model), + mapping_pattern: mappingPreviewString(matchedModel.mapping_pattern), + } + }), + } + }), + } + }) + const inferredMatches = keys.reduce( + (total, key) => total + key.matching_global_models.length, + 0, + ) + + return { + provider_id: mappingPreviewString(source.provider_id, providerId), + provider_name: mappingPreviewString(source.provider_name), + keys, + total_keys: mappingPreviewCount(source.total_keys, keys.length), + total_matches: mappingPreviewCount(source.total_matches, inferredMatches), + truncated: source.truncated === true, + truncated_keys: mappingPreviewCount(source.truncated_keys, 0), + truncated_models: mappingPreviewCount(source.truncated_models, 0), + } +} + /** * 获取 Provider 映射预览 */ @@ -379,6 +469,6 @@ export async function getProviderMappingPreview( ): Promise { return dedupedRequest(`providers:mapping-preview:${providerId}`, async () => { const response = await client.get(`/api/admin/providers/${providerId}/mapping-preview`) - return response.data + return normalizeProviderMappingPreview(response.data, providerId) }) } diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index 39efc1a28..6969c6095 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -54,6 +54,7 @@ export interface UsageRecordDetail { id: string provider?: string // 仅管理员可见 model: string + requested_reasoning_effort?: string | null reasoning_effort?: string | null service_tier?: string | null actual_service_tier?: string | null @@ -368,6 +369,7 @@ export const meApi = { has_format_conversion?: boolean | null has_fallback?: boolean | null target_model?: string | null + requested_reasoning_effort?: string | null reasoning_effort?: string | null service_tier?: string | null actual_service_tier?: string | null diff --git a/frontend/src/api/usage.ts b/frontend/src/api/usage.ts index dd31e45e6..d257d53e6 100644 --- a/frontend/src/api/usage.ts +++ b/frontend/src/api/usage.ts @@ -14,6 +14,7 @@ export interface UsageRecord { provider_id?: string // UUID provider_name?: string model: string + requested_reasoning_effort?: string | null reasoning_effort?: string | null service_tier?: string | null actual_service_tier?: string | null @@ -568,6 +569,7 @@ export const usageApi = { has_format_conversion?: boolean | null has_fallback?: boolean | null target_model?: string | null + requested_reasoning_effort?: string | null reasoning_effort?: string | null service_tier?: string | null actual_service_tier?: string | null diff --git a/frontend/src/api/users.ts b/frontend/src/api/users.ts index 10bf5419d..bc4b682d0 100644 --- a/frontend/src/api/users.ts +++ b/frontend/src/api/users.ts @@ -374,7 +374,10 @@ export const usersApi = { async listUserGroups(): Promise { const response = await apiClient.get('/api/admin/user-groups') - return response.data + return { + ...response.data, + items: Array.isArray(response.data?.items) ? response.data.items : [], + } }, async createUserGroup(payload: UpsertUserGroupRequest): Promise { @@ -417,8 +420,9 @@ export const usersApi = { }, async getUserApiKeys(userId: string): Promise { - const response = await apiClient.get<{ api_keys: ApiKey[] }>(`/api/admin/users/${userId}/api-keys`) - return response.data.api_keys + const response = await apiClient.get<{ api_keys?: ApiKey[] } | ApiKey[]>(`/api/admin/users/${userId}/api-keys`) + if (Array.isArray(response.data)) return response.data + return Array.isArray(response.data?.api_keys) ? response.data.api_keys : [] }, async getUserSessions(userId: string): Promise { diff --git a/frontend/src/features/providers/components/ProviderDetailDrawer.vue b/frontend/src/features/providers/components/ProviderDetailDrawer.vue index cffe3cb22..4cf7fd73d 100644 --- a/frontend/src/features/providers/components/ProviderDetailDrawer.vue +++ b/frontend/src/features/providers/components/ProviderDetailDrawer.vue @@ -782,7 +782,7 @@ :models="providerModels" :endpoints="endpoints" :provider-keys="providerKeys" - :loading="loadingProviderModels || loadingProviderKeys" + :loading="loadingProviderModels" @edit-model="handleEditModel" @batch-assign="handleBatchAssign" @refresh="loadEndpoints" @@ -798,7 +798,7 @@ :provider-keys="providerKeys" :models="providerModels" :mapping-preview="providerMappingPreview" - :loading="loadingProviderEndpoints || loadingProviderKeys || loadingProviderModels || loadingProviderMappingPreview" + :loading="loadingProviderMappingPreview" @refresh="handleModelMappingChanged" /> @@ -1303,8 +1303,6 @@ watch( loading.value = false } void loadSystemFormatConversionConfig() - // mapping-preview 较慢，不阻塞首屏渲染 - void loadMappingPreview() if (!hasInitialProvider) { await loadProvider() } @@ -1313,7 +1311,13 @@ watch( if (newOpen && !oldOpen) { startCountdownTimer() } - void endpointsPromise.then(() => autoRefreshQuotaInBackground()) + // 优先完成端点、密钥和模型的首屏数据，再请求计算量较大的映射预览。 + // 同时校验抽屉状态，避免关闭或切换 Provider 后启动无用请求。 + void endpointsPromise.then(() => { + if (!props.open || props.providerId !== newId) return + void loadMappingPreview() + void autoRefreshQuotaInBackground() + }) } else if (!newOpen && oldOpen) { // 使在途请求失效，避免关闭后旧响应回写 providerLoadRequestId += 1 diff --git a/frontend/src/features/providers/components/__tests__/ModelMappingTab.contract.spec.ts b/frontend/src/features/providers/components/__tests__/ModelMappingTab.contract.spec.ts new file mode 100644 index 000000000..17641ef3c --- /dev/null +++ b/frontend/src/features/providers/components/__tests__/ModelMappingTab.contract.spec.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest' +import { createSSRApp, h } from 'vue' +import { renderToString } from '@vue/server-renderer' + +import type { ProviderWithEndpointsSummary } from '@/api/endpoints' +import ModelMappingTab from '../provider-tabs/ModelMappingTab.vue' + +const provider = { + id: 'provider-demo', + name: 'Demo Provider', + provider_type: 'custom', + is_active: true, + active_keys: 0, + api_formats: [], +} as ProviderWithEndpointsSummary + +describe('ModelMappingTab response contracts', () => { + it('keeps the module visible when a legacy or malformed preview reaches the component', async () => { + const app = createSSRApp({ + render: () => h(ModelMappingTab, { + provider, + models: [], + endpoints: [], + providerKeys: [], + mappingPreview: { + message: '演示模式：该接口暂未模拟', + demo_mode: true, + }, + loading: false, + }), + }) + + const html = await renderToString(app) + + expect(html).toContain('模型映射') + expect(html).toContain('暂无模型映射') + }) +}) diff --git a/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts b/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts new file mode 100644 index 000000000..c99899f43 --- /dev/null +++ b/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts @@ -0,0 +1,32 @@ +import { readFileSync } from 'node:fs' +import { resolve } from 'node:path' +import { describe, expect, it } from 'vitest' + +const source = readFileSync( + resolve(process.cwd(), 'src/features/providers/components/ProviderDetailDrawer.vue'), + 'utf8', +) + +describe('ProviderDetailDrawer loading priorities', () => { + it('loads mapping preview after first-screen provider data', () => { + const openWatcher = source + .split('// 合并监听 providerId 和 open')[1] + ?.split('} else if (!newOpen && oldOpen)')[0] + + expect(openWatcher).toBeTruthy() + expect(openWatcher).toContain('const endpointsPromise = loadEndpoints()') + expect(openWatcher).toContain('endpointsPromise.then(() => {') + expect(openWatcher).toContain('if (!props.open || props.providerId !== newId) return') + expect(openWatcher).toContain('void loadMappingPreview()') + + const beforeEndpoints = openWatcher?.split('const endpointsPromise = loadEndpoints()')[0] + expect(beforeEndpoints).not.toContain('loadMappingPreview()') + }) + + it('keeps model and mapping loading states independent', () => { + expect(source).toContain(':loading="loadingProviderModels"') + expect(source).toContain(':loading="loadingProviderMappingPreview"') + expect(source).not.toContain(':loading="loadingProviderModels || loadingProviderKeys"') + expect(source).not.toContain(':loading="loadingProviderEndpoints || loadingProviderKeys || loadingProviderModels || loadingProviderMappingPreview"') + }) +}) diff --git a/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue b/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue index 7a0b99337..a312665b6 100644 --- a/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue +++ b/frontend/src/features/providers/components/provider-tabs/ModelMappingTab.vue @@ -540,13 +540,20 @@ const regexMappings = computed(() => { const result: CombinedMapping[] = [] const modelMap = new Map() - for (const keyInfo of aliasMappingPreview.value.keys) { - for (const gm of keyInfo.matching_global_models) { + const previewKeys = Array.isArray(aliasMappingPreview.value.keys) + ? aliasMappingPreview.value.keys + : [] + for (const keyInfo of previewKeys) { + const matchingGlobalModels = Array.isArray(keyInfo.matching_global_models) + ? keyInfo.matching_global_models + : [] + for (const gm of matchingGlobalModels) { + const matchedModels = Array.isArray(gm.matched_models) ? gm.matched_models : [] if (!modelMap.has(gm.global_model_id)) { modelMap.set(gm.global_model_id, { key: `regex-${gm.global_model_id}`, type: 'regex', - targetModelName: gm.display_name, + targetModelName: gm.display_name || gm.global_model_name || gm.global_model_id, targetModelId: gm.global_model_id, globalModelName: gm.global_model_name, mappings: [], @@ -559,7 +566,7 @@ const regexMappings = computed(() => { if (!mapping) continue // 添加 Key 信息 - const keyMatches: MappingItem[] = gm.matched_models.map(m => ({ + const keyMatches: MappingItem[] = matchedModels.map(m => ({ name: m.allowed_model, pattern: m.mapping_pattern })) @@ -572,7 +579,7 @@ const regexMappings = computed(() => { }) // 收集所有映射（去重） - for (const match of gm.matched_models) { + for (const match of matchedModels) { if (!mapping.mappings.some(m => m.name === match.allowed_model)) { mapping.mappings.push({ name: match.allowed_model, diff --git a/frontend/src/features/usage/components/ElapsedTimeText.vue b/frontend/src/features/usage/components/ElapsedTimeText.vue index 5d005162f..28d7d8391 100644 --- a/frontend/src/features/usage/components/ElapsedTimeText.vue +++ b/frontend/src/features/usage/components/ElapsedTimeText.vue @@ -22,6 +22,9 @@ const props = withDefaults(defineProps<{ const now = ref(Date.now()) const precision = computed(() => Math.max(0, props.precision)) const isActive = computed(() => props.status === 'pending' || props.status === 'streaming') +// Usage timestamps have second precision while durations have millisecond precision. +// Switching anchors can therefore introduce a sub-second phase shift at first byte. +const ACTIVE_CLOCK_TIMESTAMP_PRECISION_MS = 1000 let rafId: number | null = null @@ -72,19 +75,29 @@ const displayText = computed(() => { return `${(responseTimeMs / 1000).toFixed(precision.value)}s` } + const createdAtMs = parseCreatedAtMs(props.createdAt) + const createdAtElapsedMs = Number.isNaN(createdAtMs) + ? null + : Math.max(0, now.value - createdAtMs) + const responseTimeMs = finiteNonNegativeMs(props.responseTimeMs) const updatedAtMs = parseCreatedAtMs(props.responseTimeUpdatedAt) if (responseTimeMs != null && !Number.isNaN(updatedAtMs)) { const elapsedSinceUpdateMs = Math.max(0, now.value - updatedAtMs) - return `${((responseTimeMs + elapsedSinceUpdateMs) / 1000).toFixed(precision.value)}s` + const responseElapsedMs = responseTimeMs + elapsedSinceUpdateMs + + // When both clocks differ only by timestamp truncation, keep the original + // created-at clock so the first-byte snapshot cannot make total time pause + // or move backwards. A larger difference is a real calibration signal + // (for example an audit row created before execution) and remains authoritative. + if (createdAtElapsedMs != null && + Math.abs(responseElapsedMs - createdAtElapsedMs) <= ACTIVE_CLOCK_TIMESTAMP_PRECISION_MS) { + return `${(createdAtElapsedMs / 1000).toFixed(precision.value)}s` + } + return `${(responseElapsedMs / 1000).toFixed(precision.value)}s` } - if (!props.createdAt) return '-' - - const createdAtMs = parseCreatedAtMs(props.createdAt) - if (Number.isNaN(createdAtMs)) return '-' - - const elapsedMs = Math.max(0, now.value - createdAtMs) - return `${(elapsedMs / 1000).toFixed(precision.value)}s` + if (createdAtElapsedMs == null) return '-' + return `${(createdAtElapsedMs / 1000).toFixed(precision.value)}s` }) diff --git a/frontend/src/features/usage/components/HorizontalRequestTimeline.vue b/frontend/src/features/usage/components/HorizontalRequestTimeline.vue index 85b1e2762..6e41929fb 100644 --- a/frontend/src/features/usage/components/HorizontalRequestTimeline.vue +++ b/frontend/src/features/usage/components/HorizontalRequestTimeline.vue @@ -462,7 +462,7 @@ - +
+
+
+
@@ -1227,7 +1239,7 @@ const normalizeUpstreamResponseDisplay = (value: unknown): Record | null + diagnostic: Record | null } | null>(() => { const attempt = currentAttempt.value if (!attempt || attempt.status !== 'failed') return null @@ -1667,10 +1680,20 @@ const currentAttemptRequestError = computed<{ rawMessage, ) : null - const upstreamResponseWithDiagnostic = diagnostic - ? { ...(upstreamResponseDisplay ?? {}), diagnostic } - : upstreamResponseDisplay - if (!message && statusCode == null && !upstreamResponseWithDiagnostic) return null + const upstreamResponseData: Record = {} + const responseHeader = upstreamResponseDisplay?.headers + const responseBody = upstreamResponseDisplay?.body + if (hasRenderableValue(responseHeader)) upstreamResponseData.header = responseHeader + if (hasRenderableValue(responseBody)) upstreamResponseData.body = responseBody + const response = Object.keys(upstreamResponseData).length > 0 + ? upstreamResponseData + : null + if ( + !message + && statusCode == null + && !response + && !diagnostic + ) return null const showMessage = shouldShowAttemptMessageWithUpstreamResponse( rawMessage || fallbackType, upstreamResponseDisplay, @@ -1679,7 +1702,8 @@ const currentAttemptRequestError = computed<{ return { message: showMessage ? (message || '未知错误') : '', statusCode, - upstreamResponse: upstreamResponseWithDiagnostic, + upstreamResponse: response, + diagnostic, } }) diff --git a/frontend/src/features/usage/components/RequestDetailDrawer.vue b/frontend/src/features/usage/components/RequestDetailDrawer.vue index 5f768f36b..2cb8a36b4 100644 --- a/frontend/src/features/usage/components/RequestDetailDrawer.vue +++ b/frontend/src/features/usage/components/RequestDetailDrawer.vue @@ -23,23 +23,21 @@

请求详情

-
- {{ detail?.model || '-' }} - + +
+ -
- - -
-
- -
-
-
-

- {{ failureNotice.title }} -

- - 调度阶段 - -
-

- {{ failureNotice.message }} -

-
- - {{ item }} - -
-
-
- -
@@ -253,8 +210,6 @@ v-if="hasServiceTierFacts || processingTierPriceMultiplier !== null" class="mt-3" :requested="serviceTierFacts.requested" - :actual="serviceTierFacts.actual" - :billing="serviceTierFacts.billing" :price-multiplier="processingTierPriceMultiplier" />
@@ -886,6 +841,7 @@ import TabsContent from '@/components/ui/tabs-content.vue' import { AlertTriangle, Check, Columns2, RefreshCw, X, Monitor, Server, MessageSquareText, Code2, Terminal, Play } from 'lucide-vue-next' import { dashboardApi, type RequestDetail } from '@/api/dashboard' import type { ImageProgress, RequestTrace } from '@/api/requestTrace' +import type { UsageRecord } from '../types' import { formatApiFormat } from '@/api/endpoints/types/api-format' import { formatByteSize, @@ -906,7 +862,11 @@ import { resolveDisplayRequestStatus, resolveUsageStreamLabelSegments, } from '../utils/status' -import { resolveRequestFailureNotice } from '../utils/errorNotice' +import { isCyberPolicyError } from '../utils/cyberError' +import { + mergeUsageRecordErrorMessage, + parseUsageTimestampMs, +} from '../utils/recordSync' import { formatPricePerMillion, resolveProcessingTierPriceMultiplier, @@ -922,7 +882,11 @@ import ConversationView from './RequestDetailDrawer/ConversationView.vue' import HorizontalRequestTimeline from './HorizontalRequestTimeline.vue' import ReplayDialog from './ReplayDialog.vue' import ServiceTierFacts from './ServiceTierFacts.vue' -import { hasServiceTierFact, resolveServiceTierFacts } from '../utils/service-tier' +import UsageModelDisplay from './UsageModelDisplay.vue' +import { + hasServiceTierFact, + resolveServiceTierFacts, +} from '../utils/service-tier' // 对话解析器 import { @@ -937,6 +901,7 @@ type RequestStateStatus = 'pending' | 'streaming' | 'completed' | 'failed' | 'ca const props = defineProps<{ isOpen: boolean requestId: string | null + summaryRecord?: UsageRecord | null }>() const emit = defineEmits<{ @@ -966,11 +931,13 @@ const emit = defineEmits<{ endpointApiFormat?: string | null hasFormatConversion?: boolean | null targetModel?: string | null + requestedReasoningEffort?: string | null reasoningEffort?: string | null serviceTier?: string | null actualServiceTier?: string | null imageProgress?: ImageProgress | null errorMessage?: string | null + updatedAt?: string | null }] }>() @@ -1098,6 +1065,114 @@ function resolveRequestStateStatusFromDetail(nextDetail: Pick([ + 'target_model', + 'reasoning_effort', + 'service_tier', + 'actual_service_tier', +]) + +let modelSnapshotRevision = 0 +const summaryModelRevision = ref(0) +const detailModelRevision = ref(0) + +function usageSnapshotUpdatedAtMs( + source: UsageRecord | RequestDetail | null | undefined, +): number | null { + const value = source?.updated_at + if (typeof value !== 'string' || !value.trim()) return null + return parseUsageTimestampMs(value) +} + +function summaryNullIsNewerForProviderField( + field: HeaderModelTextField, + nextDetail: RequestDetail | null | undefined, +): boolean { + if (!FINAL_PROVIDER_HEADER_FIELDS.has(field) || !props.summaryRecord) return false + + const summaryUpdatedAt = usageSnapshotUpdatedAtMs(props.summaryRecord) + const detailUpdatedAt = usageSnapshotUpdatedAtMs(nextDetail) + if (summaryUpdatedAt != null && detailUpdatedAt != null && summaryUpdatedAt !== detailUpdatedAt) { + return summaryUpdatedAt > detailUpdatedAt + } + + if (summaryModelRevision.value > detailModelRevision.value) return true + + // A terminal list row is a complete final-provider snapshot. When no + // comparable timestamps exist, its explicit null must beat a cached detail + // from an earlier candidate. Non-terminal rows may still be filled by a + // detail request that completed after the lightweight list response. + return ['completed', 'failed', 'cancelled'].includes(props.summaryRecord.status ?? '') +} + +function readHeaderModelTextField( + source: UsageRecord | RequestDetail | null | undefined, + field: HeaderModelTextField, +): { resolved: boolean, value: string | null } { + if (!source || !Object.prototype.hasOwnProperty.call(source, field)) { + return { resolved: false, value: null } + } + + const value = (source as unknown as Record)[field] + if (value === null) return { resolved: true, value: null } + if (typeof value !== 'string') return { resolved: false, value: null } + + const normalized = value.trim() + return { resolved: true, value: normalized || null } +} + +function resolveHeaderModelTextField( + field: HeaderModelTextField, + nextDetail: RequestDetail | null | undefined, +): string | null | undefined { + // Prefer a populated list/active fact so sparse detail cannot make the header + // flicker. A summary null is often only a lightweight-contract placeholder, + // though, so a later populated detail is still useful. Final-provider stale + // facts are cleared when full list/active snapshots merge into the summary. + const summaryValue = readHeaderModelTextField(props.summaryRecord, field) + if (summaryValue.value) return summaryValue.value + + const detailValue = readHeaderModelTextField(nextDetail, field) + if (detailValue.value) { + if ( + summaryValue.resolved + && summaryValue.value === null + && summaryNullIsNewerForProviderField(field, nextDetail) + ) return null + return detailValue.value + } + + return summaryValue.resolved || detailValue.resolved ? null : undefined +} + +watch( + () => [ + props.requestId, + props.summaryRecord?.status, + props.summaryRecord?.updated_at, + props.summaryRecord?.model, + props.summaryRecord?.target_model, + props.summaryRecord?.model_version, + props.summaryRecord?.requested_reasoning_effort, + props.summaryRecord?.reasoning_effort, + props.summaryRecord?.service_tier, + props.summaryRecord?.actual_service_tier, + ], + () => { + summaryModelRevision.value = ++modelSnapshotRevision + }, + { immediate: true }, +) + function detailTotalCost(nextDetail: RequestDetail): number | null { const structuredCost = typeof nextDetail.cost === 'object' ? nextDetail.cost?.total : null const totalCost = toNumber(nextDetail.total_cost) @@ -1124,6 +1199,15 @@ function emitDetailRequestState(nextDetail: RequestDetail) { const id = props.requestId if (!id) return + const targetModel = resolveHeaderModelTextField('target_model', nextDetail) + const requestedReasoningEffort = resolveHeaderModelTextField( + 'requested_reasoning_effort', + nextDetail, + ) + const reasoningEffort = resolveHeaderModelTextField('reasoning_effort', nextDetail) + const serviceTier = resolveHeaderModelTextField('service_tier', nextDetail) + const actualServiceTier = resolveHeaderModelTextField('actual_service_tier', nextDetail) + emit('requestState', { id, requestId: nextDetail.request_id || nextDetail.id || null, @@ -1148,11 +1232,13 @@ function emitDetailRequestState(nextDetail: RequestDetail) { apiFormat: nextDetail.api_format ?? null, endpointApiFormat: nextDetail.endpoint_api_format ?? null, hasFormatConversion: nextDetail.has_format_conversion ?? null, - targetModel: nextDetail.target_model ?? null, - reasoningEffort: nextDetail.reasoning_effort ?? null, - serviceTier: nextDetail.service_tier ?? null, - actualServiceTier: nextDetail.actual_service_tier ?? null, + ...(targetModel ? { targetModel } : {}), + ...(requestedReasoningEffort ? { requestedReasoningEffort } : {}), + ...(reasoningEffort ? { reasoningEffort } : {}), + ...(serviceTier ? { serviceTier } : {}), + ...(actualServiceTier ? { actualServiceTier } : {}), errorMessage: nextDetail.error_message ?? undefined, + updatedAt: nextDetail.updated_at ?? undefined, }) } @@ -1358,10 +1444,106 @@ const metadataPanelData = computed | null>(() => { : null }) -const failureNotice = computed(() => resolveRequestFailureNotice(detail.value)) +const detailForCurrentRequest = computed(() => ( + detailMatchesRequestId(detail.value, props.requestId) ? detail.value : null +)) -const serviceTierFacts = computed(() => resolveServiceTierFacts(detail.value)) +type AuthoritativeErrorSource = 'summary' | 'detail' | null + +function isTerminalRequestState(status: RequestStateStatus | undefined): boolean { + return status === 'completed' || status === 'failed' || status === 'cancelled' +} + +function isSuccessfulTerminalRequestState(status: RequestStateStatus | undefined): boolean { + return status === 'completed' || status === 'cancelled' +} + +const authoritativeErrorSource = computed(() => { + const summary = props.summaryRecord + const currentDetail = detailForCurrentRequest.value + if (!summary || !currentDetail) return null + + const summaryStatus = resolveRequestStateStatus( + summary.status, + summary.status_code, + summary.error_message, + ) + const detailStatus = resolveRequestStateStatusFromDetail(currentDetail) + const summaryUpdatedAtMs = usageSnapshotUpdatedAtMs(summary) + const detailUpdatedAtMs = usageSnapshotUpdatedAtMs(currentDetail) + + if (summaryUpdatedAtMs != null && detailUpdatedAtMs != null && + summaryUpdatedAtMs !== detailUpdatedAtMs) { + if (detailUpdatedAtMs > summaryUpdatedAtMs && isTerminalRequestState(detailStatus)) { + return 'detail' + } + if (summaryUpdatedAtMs > detailUpdatedAtMs && isTerminalRequestState(summaryStatus)) { + return 'summary' + } + } + + // Without a comparable timestamp, a successful/cancelled terminal snapshot + // still has to clear a failure from the other source. Generic failed detail + // remains non-authoritative so opening the drawer cannot flash away a Cyber + // refusal already resolved by the list. + const detailSucceeded = isSuccessfulTerminalRequestState(detailStatus) + const summarySucceeded = isSuccessfulTerminalRequestState(summaryStatus) + if (detailSucceeded && !summarySucceeded) return 'detail' + if (summarySucceeded && !detailSucceeded) return 'summary' + if (detailSucceeded && summarySucceeded) { + return detailModelRevision.value >= summaryModelRevision.value ? 'detail' : 'summary' + } + + return null +}) + +const headerModelRecord = computed(() => { + const summary = props.summaryRecord + const currentDetail = detailForCurrentRequest.value + if (!summary && !currentDetail) return null + + const authoritativeSource = authoritativeErrorSource.value + const errorMessage = authoritativeSource === 'summary' + ? mergeUsageRecordErrorMessage(undefined, summary?.error_message, { authoritative: true }) + : mergeUsageRecordErrorMessage( + summary?.error_message, + currentDetail?.error_message, + { authoritative: authoritativeSource === 'detail' }, + ) + + return { + model: resolveHeaderModelTextField('model', currentDetail) ?? '-', + target_model: resolveHeaderModelTextField('target_model', currentDetail), + model_version: resolveHeaderModelTextField('model_version', currentDetail), + requested_reasoning_effort: resolveHeaderModelTextField( + 'requested_reasoning_effort', + currentDetail, + ), + reasoning_effort: resolveHeaderModelTextField('reasoning_effort', currentDetail), + service_tier: resolveHeaderModelTextField('service_tier', currentDetail), + error_message: errorMessage, + } +}) +const serviceTierFacts = computed(() => resolveServiceTierFacts(headerModelRecord.value)) const hasServiceTierFacts = computed(() => hasServiceTierFact(serviceTierFacts.value)) +const detailCyberPolicyError = computed(() => { + const summaryError = props.summaryRecord?.error_message + const currentDetail = detailForCurrentRequest.value + const detailErrors = [ + currentDetail?.error_message, + currentDetail?.upstream_error, + currentDetail?.failure_summary, + currentDetail?.response_body, + ] + + if (authoritativeErrorSource.value === 'summary') { + return isCyberPolicyError(summaryError) + } + if (authoritativeErrorSource.value === 'detail') { + return isCyberPolicyError(detailErrors) + } + return isCyberPolicyError([summaryError, ...detailErrors]) +}) const processingTierPriceMultiplier = computed(() => ( resolveProcessingTierPriceMultiplier(detail.value) )) @@ -2257,15 +2439,12 @@ const visibleTabs = computed(() => { }) }) -watch(() => props.requestId, async (newId) => { - if (newId && props.isOpen) { - await loadDetail(newId) - } -}) - -watch(() => props.isOpen, async (isOpen) => { - if (isOpen && props.requestId) { - await loadDetail(props.requestId) +watch([() => props.isOpen, () => props.requestId], async ([isOpen, requestId]) => { + if (isOpen && requestId) { + if (!detailMatchesRequestId(detail.value, requestId)) { + detail.value = null + } + await loadDetail(requestId) } else if (!isOpen) { stopAutoRefresh() showTimeline.value = false @@ -2276,6 +2455,14 @@ watch(() => props.isOpen, async (isOpen) => { } }) +function detailMatchesRequestId( + candidate: RequestDetail | null | undefined, + requestId: string | null | undefined, +): boolean { + if (!candidate || !requestId) return false + return candidate.id === requestId || candidate.request_id === requestId +} + async function ensureBodyContentLoaded() { if (!props.requestId || !detail.value) return @@ -2334,6 +2521,9 @@ async function loadDetail(id: string, silent = false) { const requestId = ++loadDetailRequestId loadDetailInFlight = true if (!silent) { + if (!detailMatchesRequestId(detail.value, id)) { + detail.value = null + } loading.value = true historicalPricing.value = null timelineLoaded.value = false @@ -2371,6 +2561,7 @@ async function loadDetail(id: string, silent = false) { error_flow: response.error_flow, scheduling_failure: response.scheduling_failure, } + detailModelRevision.value = ++modelSnapshotRevision detail.value = nextDetail bodiesLoadedForRequestId.value = sameRequest ? bodiesLoadedForRequestId.value : null emitDetailRequestState(nextDetail) diff --git a/frontend/src/features/usage/components/ServiceTierFacts.vue b/frontend/src/features/usage/components/ServiceTierFacts.vue index c3a33d1b8..0aa78dc80 100644 --- a/frontend/src/features/usage/components/ServiceTierFacts.vue +++ b/frontend/src/features/usage/components/ServiceTierFacts.vue @@ -1,12 +1,12 @@