use std::collections::BTreeMap; use aether_contracts::{ExecutionPlan, ExecutionTelemetry}; use aether_data_contracts::repository::usage::{UpsertUsageRecord, UsageBodyCaptureState}; use aether_data_contracts::DataLayerError; use base64::Engine as _; use serde_json::{json, Map, Value}; use crate::body_capture::{ append_runtime_body_capture_metadata, build_payload_body_capture_metadata, build_plan_body_capture_metadata, build_runtime_body_capture_states, decoded_base64_len_hint, RuntimeBodyCaptureMetadataInput, }; use crate::request_metadata::{ build_usage_request_metadata_seed, merge_usage_request_metadata, merge_usage_request_metadata_owned, sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, }; use crate::{ map_usage_from_response, GatewayStreamReportRequest, GatewaySyncReportRequest, StandardizedUsage, UsageEvent, UsageEventData, UsageEventType, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum UsageLifecycleState { Pending, Streaming, } #[derive(Debug, Clone, Default, PartialEq, Eq)] struct UsageRoutingSeed { candidate_id: Option, candidate_index: Option, key_name: Option, planner_kind: Option, route_family: Option, route_kind: Option, execution_path: Option, local_execution_runtime_miss_reason: Option, } #[derive(Debug, Clone, Default, PartialEq, Eq)] struct UsageBodyRefsSeed { request_body_ref: Option, provider_request_body_ref: Option, response_body_ref: Option, client_response_body_ref: Option, } #[derive(Debug, Clone, Default, PartialEq, Eq)] struct UsageBodyStatesSeed { request_body_state: Option, provider_request_body_state: Option, response_body_state: Option, client_response_body_state: Option, } #[derive(Debug, Clone)] struct RuntimeRequestCaptureSeed { request_body: Option, request_body_ref: Option, provider_request: Option, provider_request_body_ref: Option, body_states: UsageBodyStatesSeed, } #[derive(Debug, Clone, PartialEq)] pub struct LifecycleUsageSeed { pub request_id: String, pub user_id: Option, pub api_key_id: Option, pub username: Option, pub api_key_name: Option, pub provider_name: String, pub model: String, pub target_model: Option, pub model_id: Option, pub global_model_id: Option, pub provider_id: Option, pub provider_endpoint_id: Option, pub provider_api_key_id: Option, pub request_type: String, pub api_format: Option, pub api_family: Option, pub endpoint_kind: Option, pub endpoint_api_format: Option, pub provider_api_family: Option, pub provider_endpoint_kind: Option, pub has_format_conversion: Option, pub is_stream: bool, routing: UsageRoutingSeed, body_states: UsageBodyStatesSeed, pub request_metadata: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum UsageTerminalState { Completed, Failed, Cancelled, } #[derive(Debug, Clone)] pub struct TerminalUsageContextSeed { pub client_contract: String, pub provider_contract: String, pub request_id: String, pub user_id: Option, pub api_key_id: Option, pub username: Option, pub api_key_name: Option, pub provider_name: String, pub model: String, pub target_model: Option, pub model_id: Option, pub global_model_id: Option, pub provider_id: Option, pub provider_endpoint_id: Option, pub provider_api_key_id: Option, pub request_type: String, pub has_format_conversion: bool, pub is_stream: bool, pub request_headers: Option, pub request_body: Option, pub provider_request_headers: Option, pub provider_request: Option, body_refs: UsageBodyRefsSeed, body_states: UsageBodyStatesSeed, routing: UsageRoutingSeed, pub request_metadata: Option, } #[derive(Debug, Clone)] pub struct SyncTerminalUsagePayloadSeed { pub report_kind: String, pub status_code: u16, pub response_time_ms: Option, pub first_byte_time_ms: Option, pub provider_response_headers: Option, pub client_response_headers: Option, pub provider_response_full: Option, pub provider_response_body_state: Option, pub client_response: Option, pub client_response_body_state: Option, pub capture_metadata: Option, } #[derive(Debug, Clone)] pub struct StreamTerminalUsagePayloadSeed { pub report_kind: String, pub status_code: u16, pub response_time_ms: Option, pub first_byte_time_ms: Option, pub provider_response_headers: Option, pub client_response_headers: Option, pub provider_response_full: Option, pub provider_response_body_state: Option, pub client_response: Option, pub client_response_body_state: Option, pub standardized_usage: Option, pub capture_metadata: Option, } #[derive(Debug, Clone)] pub struct TerminalUsageSeed { pub terminal_state: UsageTerminalState, pub client_contract: String, pub provider_contract: String, pub request_id: String, pub user_id: Option, pub api_key_id: Option, pub username: Option, pub api_key_name: Option, pub provider_name: String, pub model: String, pub target_model: Option, pub model_id: Option, pub global_model_id: Option, pub provider_id: Option, pub provider_endpoint_id: Option, pub provider_api_key_id: Option, pub request_type: String, pub has_format_conversion: bool, pub is_stream: bool, pub status_code: u16, pub response_time_ms: Option, pub first_byte_time_ms: Option, pub request_headers: Option, pub request_body: Option, pub provider_request_headers: Option, pub provider_request: Option, body_refs: UsageBodyRefsSeed, body_states: UsageBodyStatesSeed, pub provider_response_headers: Option, pub provider_response: Option, pub client_response_headers: Option, pub client_response: Option, routing: UsageRoutingSeed, pub request_metadata: Option, pub audit_payload: Option, pub standardized_usage: Option, } pub type TerminalUsageOutcome = TerminalUsageSeed; struct LifecycleUsageRecordInput<'a> { seed: &'a LifecycleUsageSeed, options: LifecycleUsageRecordOptions, } struct OwnedLifecycleUsageRecordInput { seed: LifecycleUsageSeed, options: LifecycleUsageRecordOptions, } struct LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState, status_code: Option, response_time_ms: Option, first_byte_time_ms: Option, response_headers: Option, client_response_headers: Option, updated_at_unix_secs: u64, trusted_request_metadata: bool, } const MAX_USAGE_CAPTURE_DEPTH: usize = 64; const MAX_USAGE_CAPTURE_NODES: usize = 20_000; const MAX_USAGE_CAPTURE_BYTES: usize = 64 * 1024; pub fn build_lifecycle_usage_seed( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> LifecycleUsageSeed { let context = report_context.and_then(Value::as_object); let api_format = context_string(context, "client_api_format") .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))); let endpoint_api_format = context_string(context, "provider_api_format") .or_else(|| non_empty_str(Some(plan.provider_api_format.as_str()))); let provider_name = context_string(context, "provider_name") .or_else(|| non_empty_str(plan.provider_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()); let model = context_string(context, "model") .or_else(|| non_empty_str(plan.model_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()); let request_type = infer_request_type(api_format.as_deref()); let api_family = api_format .as_deref() .and_then(infer_api_family) .map(ToOwned::to_owned); let endpoint_kind = api_format .as_deref() .and_then(infer_endpoint_kind) .map(ToOwned::to_owned); let provider_api_family = endpoint_api_format .as_deref() .and_then(infer_api_family) .map(ToOwned::to_owned); let provider_endpoint_kind = endpoint_api_format .as_deref() .and_then(infer_endpoint_kind) .map(ToOwned::to_owned); LifecycleUsageSeed { request_id: plan.request_id.clone(), user_id: context_string(context, "user_id"), api_key_id: context_string(context, "api_key_id"), username: context_string(context, "username"), api_key_name: context_string(context, "api_key_name"), provider_name, model, target_model: context_string(context, "mapped_model"), model_id: context_string(context, "model_id"), global_model_id: context_string(context, "global_model_id"), provider_id: empty_to_none( context_string(context, "provider_id") .or_else(|| non_empty_str(Some(plan.provider_id.as_str()))), ), provider_endpoint_id: empty_to_none( context_string(context, "endpoint_id") .or_else(|| non_empty_str(Some(plan.endpoint_id.as_str()))), ), provider_api_key_id: empty_to_none( context_string(context, "key_id").or_else(|| non_empty_str(Some(plan.key_id.as_str()))), ), request_type, api_format, api_family, endpoint_kind, endpoint_api_format, provider_api_family, provider_endpoint_kind, has_format_conversion: context_bool(context, "needs_conversion"), is_stream: plan.stream, routing: build_runtime_routing_seed(plan, context), body_states: build_runtime_body_states_seed(plan, context), request_metadata: build_runtime_request_metadata_seed(plan, context), } } pub fn build_pending_usage_record( plan: &ExecutionPlan, report_context: Option<&Value>, updated_at_unix_secs: u64, ) -> Result { let seed = build_lifecycle_usage_seed(plan, report_context); build_lifecycle_usage_record_owned(OwnedLifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Pending, status_code: None, response_time_ms: None, first_byte_time_ms: None, response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub fn build_pending_usage_record_from_seed( seed: &LifecycleUsageSeed, updated_at_unix_secs: u64, ) -> Result { build_lifecycle_usage_record(LifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Pending, status_code: None, response_time_ms: None, first_byte_time_ms: None, response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub(crate) fn build_pending_usage_record_from_owned_seed( seed: LifecycleUsageSeed, updated_at_unix_secs: u64, ) -> Result { build_lifecycle_usage_record_owned(OwnedLifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Pending, status_code: None, response_time_ms: None, first_byte_time_ms: None, response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub fn build_streaming_usage_record( plan: &ExecutionPlan, report_context: Option<&Value>, status_code: u16, telemetry: Option<&ExecutionTelemetry>, updated_at_unix_secs: u64, ) -> Result { let seed = build_lifecycle_usage_seed(plan, report_context); build_lifecycle_usage_record_owned(OwnedLifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Streaming, status_code: Some(status_code), response_time_ms: telemetry.and_then(|value| value.elapsed_ms), first_byte_time_ms: telemetry.and_then(|value| value.ttfb_ms), response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub fn build_streaming_usage_record_from_seed( seed: &LifecycleUsageSeed, status_code: u16, telemetry: Option<&ExecutionTelemetry>, updated_at_unix_secs: u64, ) -> Result { build_lifecycle_usage_record(LifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Streaming, status_code: Some(status_code), response_time_ms: telemetry.and_then(|value| value.elapsed_ms), first_byte_time_ms: telemetry.and_then(|value| value.ttfb_ms), response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub(crate) fn build_streaming_usage_record_from_owned_seed( seed: LifecycleUsageSeed, status_code: u16, telemetry: Option, updated_at_unix_secs: u64, ) -> Result { build_lifecycle_usage_record_owned(OwnedLifecycleUsageRecordInput { seed, options: LifecycleUsageRecordOptions { lifecycle_state: UsageLifecycleState::Streaming, status_code: Some(status_code), response_time_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms), first_byte_time_ms: telemetry.as_ref().and_then(|value| value.ttfb_ms), response_headers: None, client_response_headers: None, updated_at_unix_secs, trusted_request_metadata: false, }, }) } pub fn build_sync_terminal_usage_event( plan: &ExecutionPlan, report_context: Option<&Value>, payload: &GatewaySyncReportRequest, ) -> Result { let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_sync_terminal_usage_payload_seed(payload); build_terminal_usage_event_from_seed_impl( build_sync_terminal_usage_seed(context_seed, payload_seed), true, ) } pub fn build_stream_terminal_usage_event( plan: &ExecutionPlan, report_context: Option<&Value>, payload: &GatewayStreamReportRequest, ) -> Result { let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_stream_terminal_usage_payload_seed(payload); build_terminal_usage_event_from_seed_impl( build_stream_terminal_usage_seed(context_seed, payload_seed, false), true, ) } pub fn build_sync_terminal_usage_outcome( plan: &ExecutionPlan, report_context: Option<&Value>, payload: &GatewaySyncReportRequest, ) -> TerminalUsageOutcome { let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_sync_terminal_usage_payload_seed(payload); build_sync_terminal_usage_seed(context_seed, payload_seed) } pub fn build_stream_terminal_usage_outcome( plan: &ExecutionPlan, report_context: Option<&Value>, payload: &GatewayStreamReportRequest, ) -> TerminalUsageOutcome { let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_stream_terminal_usage_payload_seed(payload); build_stream_terminal_usage_seed(context_seed, payload_seed, false) } pub fn build_terminal_usage_event_from_outcome( outcome: TerminalUsageOutcome, ) -> Result { build_terminal_usage_event_from_seed(outcome) } pub fn build_terminal_usage_event_from_seed( seed: TerminalUsageSeed, ) -> Result { build_terminal_usage_event_from_seed_impl(seed, false) } fn build_terminal_usage_event_from_seed_impl( seed: TerminalUsageSeed, trusted_request_metadata: bool, ) -> Result { let TerminalUsageSeed { terminal_state, client_contract, provider_contract, request_id, user_id, api_key_id, username, api_key_name, provider_name, model, target_model, model_id, global_model_id, provider_id, provider_endpoint_id, provider_api_key_id, request_type, has_format_conversion, is_stream, status_code, response_time_ms, first_byte_time_ms, request_headers, request_body, provider_request_headers, provider_request, body_refs, body_states, provider_response_headers, provider_response, client_response_headers, client_response, routing, request_metadata, audit_payload, standardized_usage, } = seed; let event_type = match terminal_state { UsageTerminalState::Completed => UsageEventType::Completed, UsageTerminalState::Failed => UsageEventType::Failed, UsageTerminalState::Cancelled => UsageEventType::Cancelled, }; let routing = merge_routing_seed_with_metadata_owned(routing, request_metadata.as_ref()); let body_refs = merge_body_refs_seed_with_metadata_owned(body_refs, request_metadata.as_ref()); let error_message = resolve_error_message(status_code, provider_response.as_ref(), None) .or_else(|| resolve_error_message(status_code, client_response.as_ref(), None)); let api_family = infer_api_family(&client_contract).map(ToOwned::to_owned); let endpoint_kind = infer_endpoint_kind(&client_contract).map(ToOwned::to_owned); let provider_api_family = infer_api_family(&provider_contract).map(ToOwned::to_owned); let provider_endpoint_kind = infer_endpoint_kind(&provider_contract).map(ToOwned::to_owned); let derived_standardized_usage = provider_response .as_ref() .filter(|response_body| response_body.is_object()) .map(|response_body| map_usage_from_response(response_body, provider_contract.as_str())) .filter(StandardizedUsage::has_token_signal); let standardized_usage = StandardizedUsage::choose_more_complete(standardized_usage, derived_standardized_usage); let request_metadata = if trusted_request_metadata { merge_usage_request_metadata_owned(request_metadata, audit_payload) } else { merge_usage_request_metadata(request_metadata, audit_payload) }; let mut data = UsageEventData { user_id, api_key_id, username, api_key_name, provider_name, model, target_model, model_id, global_model_id, provider_id, provider_endpoint_id, provider_api_key_id, request_type: Some(request_type), api_format: Some(client_contract), api_family, endpoint_kind, endpoint_api_format: Some(provider_contract), provider_api_family, provider_endpoint_kind, has_format_conversion: Some(has_format_conversion), is_stream: Some(is_stream), status_code: Some(status_code), error_message, error_category: resolve_error_category(status_code, event_type), response_time_ms, first_byte_time_ms, request_headers, request_body, request_body_ref: body_refs.request_body_ref, request_body_state: body_states.request_body_state, provider_request_headers, provider_request_body: provider_request, provider_request_body_ref: body_refs.provider_request_body_ref, provider_request_body_state: body_states.provider_request_body_state, response_headers: provider_response_headers, response_body: provider_response, response_body_ref: body_refs.response_body_ref, response_body_state: body_states.response_body_state, client_response_headers, client_response_body: client_response, client_response_body_ref: body_refs.client_response_body_ref, client_response_body_state: body_states.client_response_body_state, candidate_id: routing.candidate_id, key_name: routing.key_name, planner_kind: routing.planner_kind, route_family: routing.route_family, route_kind: routing.route_kind, execution_path: routing.execution_path, local_execution_runtime_miss_reason: routing.local_execution_runtime_miss_reason, request_metadata, ..UsageEventData::default() }; if let Some(usage) = standardized_usage.as_ref() { apply_standardized_usage_seed(usage, &mut data); } if data.total_tokens.is_none() { if let Some(tokens) = data .response_body .as_ref() .and_then(extract_token_counts_from_value) { data.input_tokens = Some(tokens.0); data.output_tokens = Some(tokens.1); data.total_tokens = Some(tokens.2); } } let data = if trusted_request_metadata { sanitize_usage_event_capture_fields_trusted(data) } else { sanitize_usage_event_data(data) }; Ok(UsageEvent::new(event_type, request_id, data)) } pub fn build_terminal_usage_context_seed( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> TerminalUsageContextSeed { let context = report_context.and_then(Value::as_object); let request_capture = build_runtime_request_capture_seed(plan, context); let client_contract = context_string(context, "client_contract") .or_else(|| context_string(context, "client_api_format")) .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))) .unwrap_or_default(); let provider_contract = context_string(context, "provider_contract") .or_else(|| context_string(context, "provider_api_format")) .or_else(|| non_empty_str(Some(plan.provider_api_format.as_str()))) .unwrap_or_default(); let request_type = infer_request_type(Some(client_contract.as_str())); let has_format_conversion = resolve_has_format_conversion( context, client_contract.as_str(), provider_contract.as_str(), ); TerminalUsageContextSeed { client_contract, provider_contract, has_format_conversion, request_id: plan.request_id.clone(), user_id: context_string(context, "user_id"), api_key_id: context_string(context, "api_key_id"), username: context_string(context, "username"), api_key_name: context_string(context, "api_key_name"), provider_name: context_string(context, "provider_name") .or_else(|| non_empty_str(plan.provider_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()), model: context_string(context, "model") .or_else(|| non_empty_str(plan.model_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()), target_model: context_string(context, "mapped_model"), model_id: context_string(context, "model_id"), global_model_id: context_string(context, "global_model_id"), provider_id: context_string(context, "provider_id") .or_else(|| non_empty_str(Some(plan.provider_id.as_str()))), provider_endpoint_id: context_string(context, "endpoint_id") .or_else(|| non_empty_str(Some(plan.endpoint_id.as_str()))), provider_api_key_id: context_string(context, "key_id") .or_else(|| non_empty_str(Some(plan.key_id.as_str()))), request_type, is_stream: plan.stream, routing: build_runtime_routing_seed(plan, context), request_headers: mask_sensitive_headers_in_json_value(context_value( context, "original_headers", )), request_body: request_capture.request_body, provider_request_headers: mask_sensitive_headers_in_json_value( context_usage_value(context, "provider_request_headers") .or_else(|| headers_to_json(&plan.headers)), ), provider_request: request_capture.provider_request, body_refs: UsageBodyRefsSeed { request_body_ref: request_capture.request_body_ref, provider_request_body_ref: request_capture.provider_request_body_ref, response_body_ref: context_string(context, "response_body_ref"), client_response_body_ref: context_string(context, "client_response_body_ref"), }, body_states: request_capture.body_states, request_metadata: merge_usage_request_metadata_owned( build_usage_request_metadata_seed(plan, context), build_plan_body_capture_metadata(plan.body.body_bytes_b64.as_deref()), ), } } pub fn build_sync_terminal_usage_payload_seed( payload: &GatewaySyncReportRequest, ) -> SyncTerminalUsagePayloadSeed { let upstream_is_stream = payload .report_context .as_ref() .and_then(Value::as_object) .and_then(|context| context.get("upstream_is_stream")) .and_then(Value::as_bool) .unwrap_or(false); let provider_response_full = if upstream_is_stream && payload.body_base64.is_some() { decode_body_for_storage(payload.body_base64.as_deref()) .or_else(|| payload.body_json.as_ref().cloned()) } else { payload .body_json .as_ref() .cloned() .or_else(|| decode_body_for_storage(payload.body_base64.as_deref())) }; let has_provider_response = provider_response_full.is_some(); let client_response = payload.client_body_json.as_ref().cloned(); let has_client_response = client_response.is_some(); let provider_response_body_state = Some(UsageBodyCaptureState::from_capture_parts( has_provider_response, false, false, )); let client_response_body_state = Some(UsageBodyCaptureState::from_capture_parts( has_client_response, false, false, )); let context = payload.report_context.as_ref().and_then(Value::as_object); let provider_response_headers = context_usage_value(context, "provider_response_headers") .or_else(|| headers_to_json(&payload.headers)); let client_response_headers = context_usage_value(context, "client_response_headers") .or_else(|| headers_to_json(&payload.headers)); SyncTerminalUsagePayloadSeed { report_kind: payload.report_kind.clone(), status_code: payload.status_code, response_time_ms: payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms), provider_response_headers, client_response_headers, provider_response_full, provider_response_body_state, client_response, client_response_body_state, capture_metadata: build_payload_body_capture_metadata( payload.body_base64.as_deref(), None, provider_response_body_state, client_response_body_state, ), } } pub fn build_stream_terminal_usage_payload_seed( payload: &GatewayStreamReportRequest, ) -> StreamTerminalUsagePayloadSeed { let context = payload.report_context.as_ref().and_then(Value::as_object); let provider_response_headers = context_usage_value(context, "provider_response_headers") .or_else(|| headers_to_json(&payload.headers)); let client_response_headers = headers_to_json(&payload.headers); StreamTerminalUsagePayloadSeed { report_kind: payload.report_kind.clone(), status_code: payload.status_code, response_time_ms: payload .telemetry .as_ref() .and_then(|value| value.elapsed_ms), first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms), provider_response_headers, client_response_headers, provider_response_full: decode_body_for_storage(payload.provider_body_base64.as_deref()), provider_response_body_state: payload.provider_body_state, client_response: decode_body_for_storage(payload.client_body_base64.as_deref()), client_response_body_state: payload.client_body_state, standardized_usage: payload .terminal_summary .as_ref() .and_then(|summary| summary.standardized_usage.clone()), capture_metadata: build_payload_body_capture_metadata( payload.provider_body_base64.as_deref(), payload.client_body_base64.as_deref(), payload.provider_body_state, payload.client_body_state, ), } } pub fn build_sync_terminal_usage_seed( context_seed: TerminalUsageContextSeed, payload_seed: SyncTerminalUsagePayloadSeed, ) -> TerminalUsageSeed { let SyncTerminalUsagePayloadSeed { report_kind, status_code, response_time_ms, first_byte_time_ms, provider_response_headers, client_response_headers, provider_response_full, provider_response_body_state, client_response, client_response_body_state, capture_metadata, } = payload_seed; let standardized_usage = provider_response_full .as_ref() .map(|response| map_usage_from_response(response, context_seed.provider_contract.as_str())); let terminal_state = infer_sync_terminal_state( report_kind.as_str(), status_code, provider_response_full.as_ref(), ); TerminalUsageSeed { terminal_state, client_contract: context_seed.client_contract, provider_contract: context_seed.provider_contract, request_id: context_seed.request_id, user_id: context_seed.user_id, api_key_id: context_seed.api_key_id, username: context_seed.username, api_key_name: context_seed.api_key_name, provider_name: context_seed.provider_name, model: context_seed.model, target_model: context_seed.target_model, model_id: context_seed.model_id, global_model_id: context_seed.global_model_id, provider_id: context_seed.provider_id, provider_endpoint_id: context_seed.provider_endpoint_id, provider_api_key_id: context_seed.provider_api_key_id, request_type: context_seed.request_type, has_format_conversion: context_seed.has_format_conversion, is_stream: context_seed.is_stream, status_code, response_time_ms, first_byte_time_ms, request_headers: context_seed.request_headers, request_body: context_seed.request_body, provider_request_headers: context_seed.provider_request_headers, provider_request: context_seed.provider_request, body_refs: context_seed.body_refs, body_states: UsageBodyStatesSeed { request_body_state: context_seed.body_states.request_body_state, provider_request_body_state: context_seed.body_states.provider_request_body_state, response_body_state: provider_response_body_state, client_response_body_state, }, routing: context_seed.routing, provider_response_headers, provider_response: provider_response_full, client_response_headers, client_response, request_metadata: context_seed.request_metadata, audit_payload: capture_metadata, standardized_usage, } } pub fn build_stream_terminal_usage_seed( context_seed: TerminalUsageContextSeed, payload_seed: StreamTerminalUsagePayloadSeed, cancelled: bool, ) -> TerminalUsageSeed { let StreamTerminalUsagePayloadSeed { report_kind, status_code, response_time_ms, first_byte_time_ms, provider_response_headers, client_response_headers, provider_response_full, provider_response_body_state, client_response, client_response_body_state, standardized_usage, capture_metadata, } = payload_seed; let standardized_usage = standardized_usage.or_else(|| { provider_response_full.as_ref().map(|response| { map_usage_from_response(response, context_seed.provider_contract.as_str()) }) }); let terminal_state = infer_stream_terminal_state(report_kind.as_str(), status_code, cancelled); TerminalUsageSeed { terminal_state, client_contract: context_seed.client_contract, provider_contract: context_seed.provider_contract, request_id: context_seed.request_id, user_id: context_seed.user_id, api_key_id: context_seed.api_key_id, username: context_seed.username, api_key_name: context_seed.api_key_name, provider_name: context_seed.provider_name, model: context_seed.model, target_model: context_seed.target_model, model_id: context_seed.model_id, global_model_id: context_seed.global_model_id, provider_id: context_seed.provider_id, provider_endpoint_id: context_seed.provider_endpoint_id, provider_api_key_id: context_seed.provider_api_key_id, request_type: context_seed.request_type, has_format_conversion: context_seed.has_format_conversion, is_stream: context_seed.is_stream, status_code, response_time_ms, first_byte_time_ms, request_headers: context_seed.request_headers, request_body: context_seed.request_body, provider_request_headers: context_seed.provider_request_headers, provider_request: context_seed.provider_request, body_refs: context_seed.body_refs, body_states: UsageBodyStatesSeed { request_body_state: context_seed.body_states.request_body_state, provider_request_body_state: context_seed.body_states.provider_request_body_state, response_body_state: provider_response_body_state, client_response_body_state, }, routing: context_seed.routing, provider_response_headers, provider_response: provider_response_full, client_response_headers, client_response, request_metadata: context_seed.request_metadata, audit_payload: capture_metadata, standardized_usage, } } fn infer_sync_terminal_state( report_kind: &str, status_code: u16, provider_response: Option<&Value>, ) -> UsageTerminalState { if status_code == 499 || report_kind.contains("cancel") { UsageTerminalState::Cancelled } else if !(200..300).contains(&status_code) || provider_response .and_then(|value| value.get("error")) .is_some_and(|value| !value.is_null()) { UsageTerminalState::Failed } else { UsageTerminalState::Completed } } fn infer_stream_terminal_state( report_kind: &str, status_code: u16, cancelled: bool, ) -> UsageTerminalState { if cancelled || status_code == 499 || report_kind.contains("cancel") { UsageTerminalState::Cancelled } else if !(200..300).contains(&status_code) { UsageTerminalState::Failed } else { UsageTerminalState::Completed } } fn resolve_has_format_conversion( context: Option<&Map>, client_contract: &str, provider_contract: &str, ) -> bool { match context_string(context, "conversion_mode").as_deref() { Some("none") => false, Some("request_only" | "response_only" | "bidirectional") => true, _ if context_bool(context, "needs_conversion").unwrap_or(false) => true, _ if client_contract != provider_contract => true, _ => false, } } fn build_lifecycle_usage_record( input: LifecycleUsageRecordInput<'_>, ) -> Result { let LifecycleUsageRecordInput { seed, options } = input; build_lifecycle_usage_record_impl(seed, options) } fn build_lifecycle_usage_record_owned( input: OwnedLifecycleUsageRecordInput, ) -> Result { let OwnedLifecycleUsageRecordInput { seed, options } = input; let LifecycleUsageRecordOptions { lifecycle_state, status_code, response_time_ms, first_byte_time_ms, response_headers, client_response_headers, updated_at_unix_secs, trusted_request_metadata, } = options; let (status, billing_status) = lifecycle_status_and_billing(lifecycle_state); let LifecycleUsageSeed { request_id, user_id, api_key_id, username, api_key_name, 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, routing, body_states, request_metadata, .. } = seed; let routing = merge_routing_seed_with_metadata_owned(routing, request_metadata.as_ref()); let body_refs = merge_body_refs_seed_with_metadata_owned( UsageBodyRefsSeed::default(), request_metadata.as_ref(), ); let request_metadata = if trusted_request_metadata { request_metadata } else { sanitize_usage_request_metadata(request_metadata) }; Ok(UpsertUsageRecord { request_id, user_id, api_key_id, username, api_key_name, provider_name, model, target_model, provider_id, provider_endpoint_id, provider_api_key_id, request_type: Some(request_type), api_format, api_family, endpoint_kind, endpoint_api_format, provider_api_family, provider_endpoint_kind, has_format_conversion, is_stream: Some(is_stream), input_tokens: None, output_tokens: None, total_tokens: None, 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, error_message: None, error_category: None, response_time_ms, first_byte_time_ms, status: status.to_string(), billing_status: billing_status.to_string(), request_headers: None, request_body: None, request_body_ref: body_refs.request_body_ref, request_body_state: body_states.request_body_state, provider_request_headers: None, provider_request_body: None, provider_request_body_ref: body_refs.provider_request_body_ref, provider_request_body_state: body_states.provider_request_body_state, response_headers: sanitize_usage_header_capture(response_headers), response_body: None, response_body_ref: body_refs.response_body_ref, response_body_state: Some(UsageBodyCaptureState::None), client_response_headers: sanitize_usage_header_capture(client_response_headers), client_response_body: None, client_response_body_ref: body_refs.client_response_body_ref, client_response_body_state: Some(UsageBodyCaptureState::None), candidate_id: routing.candidate_id, candidate_index: routing.candidate_index, key_name: routing.key_name, planner_kind: routing.planner_kind, route_family: routing.route_family, route_kind: routing.route_kind, execution_path: routing.execution_path, local_execution_runtime_miss_reason: routing.local_execution_runtime_miss_reason, request_metadata, finalized_at_unix_secs: None, created_at_unix_ms: Some(updated_at_unix_secs), updated_at_unix_secs, }) } fn build_lifecycle_usage_record_impl( seed: &LifecycleUsageSeed, options: LifecycleUsageRecordOptions, ) -> Result { let LifecycleUsageRecordOptions { lifecycle_state, status_code, response_time_ms, first_byte_time_ms, response_headers, client_response_headers, updated_at_unix_secs, trusted_request_metadata, } = options; let (status, billing_status) = lifecycle_status_and_billing(lifecycle_state); let routing = merge_routing_seed_with_metadata(&seed.routing, seed.request_metadata.as_ref()); let body_refs = merge_body_refs_seed_with_metadata( &UsageBodyRefsSeed::default(), seed.request_metadata.as_ref(), ); let request_metadata = if trusted_request_metadata { seed.request_metadata.clone() } else { sanitize_usage_request_metadata_ref(seed.request_metadata.as_ref()) }; Ok(UpsertUsageRecord { request_id: seed.request_id.clone(), user_id: seed.user_id.clone(), api_key_id: seed.api_key_id.clone(), username: seed.username.clone(), api_key_name: seed.api_key_name.clone(), provider_name: seed.provider_name.clone(), model: seed.model.clone(), target_model: seed.target_model.clone(), provider_id: seed.provider_id.clone(), provider_endpoint_id: seed.provider_endpoint_id.clone(), provider_api_key_id: seed.provider_api_key_id.clone(), request_type: Some(seed.request_type.clone()), api_format: seed.api_format.clone(), api_family: seed.api_family.clone(), endpoint_kind: seed.endpoint_kind.clone(), endpoint_api_format: seed.endpoint_api_format.clone(), provider_api_family: seed.provider_api_family.clone(), provider_endpoint_kind: seed.provider_endpoint_kind.clone(), has_format_conversion: seed.has_format_conversion, is_stream: Some(seed.is_stream), input_tokens: None, output_tokens: None, total_tokens: None, 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, error_message: None, error_category: None, response_time_ms, first_byte_time_ms, status: status.to_string(), billing_status: billing_status.to_string(), request_headers: None, request_body: None, request_body_ref: body_refs.request_body_ref, request_body_state: seed.body_states.request_body_state, provider_request_headers: None, provider_request_body: None, provider_request_body_ref: body_refs.provider_request_body_ref, provider_request_body_state: seed.body_states.provider_request_body_state, response_headers: sanitize_usage_header_capture(response_headers), response_body: None, response_body_ref: body_refs.response_body_ref, response_body_state: Some(UsageBodyCaptureState::None), client_response_headers: sanitize_usage_header_capture(client_response_headers), client_response_body: None, client_response_body_ref: body_refs.client_response_body_ref, client_response_body_state: Some(UsageBodyCaptureState::None), candidate_id: routing.candidate_id, candidate_index: routing.candidate_index, key_name: routing.key_name, planner_kind: routing.planner_kind, route_family: routing.route_family, route_kind: routing.route_kind, execution_path: routing.execution_path, local_execution_runtime_miss_reason: routing.local_execution_runtime_miss_reason, request_metadata, finalized_at_unix_secs: None, created_at_unix_ms: Some(updated_at_unix_secs), updated_at_unix_secs, }) } pub fn build_usage_event_data_seed( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> UsageEventData { build_usage_event_data_seed_with_detail(plan, report_context) } fn build_usage_event_data_seed_with_detail( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> UsageEventData { let context = report_context.and_then(Value::as_object); let routing = build_runtime_routing_seed(plan, context); let request_capture = build_runtime_request_capture_seed(plan, context); let api_format = context_string(context, "client_api_format") .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))); let endpoint_api_format = context_string(context, "provider_api_format") .or_else(|| non_empty_str(Some(plan.provider_api_format.as_str()))); let model = context_string(context, "model") .or_else(|| non_empty_str(plan.model_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()); let provider_name = context_string(context, "provider_name") .or_else(|| non_empty_str(plan.provider_name.as_deref())) .unwrap_or_else(|| "unknown".to_string()); let request_type = Some(infer_request_type(api_format.as_deref())); let api_family = api_format .as_deref() .and_then(infer_api_family) .map(ToOwned::to_owned); let endpoint_kind = api_format .as_deref() .and_then(infer_endpoint_kind) .map(ToOwned::to_owned); let provider_api_family = endpoint_api_format .as_deref() .and_then(infer_api_family) .map(ToOwned::to_owned); let provider_endpoint_kind = endpoint_api_format .as_deref() .and_then(infer_endpoint_kind) .map(ToOwned::to_owned); let request_metadata = build_runtime_request_metadata_seed_from_parts( context, request_capture.request_body.is_some(), request_capture.request_body_ref.as_deref(), request_capture.provider_request.is_some(), request_capture.provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.as_deref(), ); sanitize_usage_event_data(UsageEventData { user_id: context_string(context, "user_id"), api_key_id: context_string(context, "api_key_id"), username: context_string(context, "username"), api_key_name: context_string(context, "api_key_name"), provider_name, model, target_model: context_string(context, "mapped_model"), model_id: context_string(context, "model_id"), global_model_id: context_string(context, "global_model_id"), provider_id: context_string(context, "provider_id") .or_else(|| non_empty_str(Some(plan.provider_id.as_str()))), provider_endpoint_id: context_string(context, "endpoint_id") .or_else(|| non_empty_str(Some(plan.endpoint_id.as_str()))), provider_api_key_id: context_string(context, "key_id") .or_else(|| non_empty_str(Some(plan.key_id.as_str()))), request_type, api_format, api_family, endpoint_kind, endpoint_api_format, provider_api_family, provider_endpoint_kind, has_format_conversion: context_bool(context, "needs_conversion"), is_stream: Some(plan.stream), request_headers: context_usage_value(context, "original_headers"), request_body: request_capture.request_body, request_body_ref: request_capture.request_body_ref, request_body_state: request_capture.body_states.request_body_state, provider_request_headers: context_usage_value(context, "provider_request_headers") .or_else(|| headers_to_json(&plan.headers)), provider_request_body: request_capture.provider_request, provider_request_body_ref: request_capture.provider_request_body_ref, provider_request_body_state: request_capture.body_states.provider_request_body_state, response_body_ref: context_string(context, "response_body_ref"), response_body_state: request_capture.body_states.response_body_state, client_response_body_ref: context_string(context, "client_response_body_ref"), client_response_body_state: request_capture.body_states.client_response_body_state, candidate_id: routing.candidate_id, candidate_index: routing.candidate_index, key_name: routing.key_name, planner_kind: routing.planner_kind, route_family: routing.route_family, route_kind: routing.route_kind, execution_path: routing.execution_path, local_execution_runtime_miss_reason: routing.local_execution_runtime_miss_reason, request_metadata, ..UsageEventData::default() }) } fn lifecycle_status_and_billing(state: UsageLifecycleState) -> (&'static str, &'static str) { match state { UsageLifecycleState::Pending => ("pending", "pending"), UsageLifecycleState::Streaming => ("streaming", "pending"), } } fn build_runtime_routing_seed( plan: &ExecutionPlan, context: Option<&Map>, ) -> UsageRoutingSeed { UsageRoutingSeed { candidate_id: context_string(context, "candidate_id") .or_else(|| non_empty_str(plan.candidate_id.as_deref())), candidate_index: context_u64(context, "candidate_index"), key_name: context_string(context, "key_name"), planner_kind: context_string(context, "planner_kind"), route_family: context_string(context, "route_family"), route_kind: context_string(context, "route_kind"), execution_path: context_string(context, "execution_path"), local_execution_runtime_miss_reason: context_string( context, "local_execution_runtime_miss_reason", ), } } fn routing_string_from_metadata(value: Option<&Value>, key: &str) -> Option { value .and_then(Value::as_object) .and_then(|object| object.get(key)) .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) } fn merge_routing_seed_with_metadata( routing: &UsageRoutingSeed, metadata: Option<&Value>, ) -> UsageRoutingSeed { UsageRoutingSeed { candidate_id: routing .candidate_id .clone() .or_else(|| routing_string_from_metadata(metadata, "candidate_id")), candidate_index: routing .candidate_index .or_else(|| routing_u64_from_metadata(metadata, "candidate_index")), key_name: routing .key_name .clone() .or_else(|| routing_string_from_metadata(metadata, "key_name")), planner_kind: routing .planner_kind .clone() .or_else(|| routing_string_from_metadata(metadata, "planner_kind")), route_family: routing .route_family .clone() .or_else(|| routing_string_from_metadata(metadata, "route_family")), route_kind: routing .route_kind .clone() .or_else(|| routing_string_from_metadata(metadata, "route_kind")), execution_path: routing .execution_path .clone() .or_else(|| routing_string_from_metadata(metadata, "execution_path")), local_execution_runtime_miss_reason: routing .local_execution_runtime_miss_reason .clone() .or_else(|| { routing_string_from_metadata(metadata, "local_execution_runtime_miss_reason") }), } } fn merge_routing_seed_with_metadata_owned( routing: UsageRoutingSeed, metadata: Option<&Value>, ) -> UsageRoutingSeed { UsageRoutingSeed { candidate_id: routing .candidate_id .or_else(|| routing_string_from_metadata(metadata, "candidate_id")), candidate_index: routing .candidate_index .or_else(|| routing_u64_from_metadata(metadata, "candidate_index")), key_name: routing .key_name .or_else(|| routing_string_from_metadata(metadata, "key_name")), planner_kind: routing .planner_kind .or_else(|| routing_string_from_metadata(metadata, "planner_kind")), route_family: routing .route_family .or_else(|| routing_string_from_metadata(metadata, "route_family")), route_kind: routing .route_kind .or_else(|| routing_string_from_metadata(metadata, "route_kind")), execution_path: routing .execution_path .or_else(|| routing_string_from_metadata(metadata, "execution_path")), local_execution_runtime_miss_reason: routing.local_execution_runtime_miss_reason.or_else( || routing_string_from_metadata(metadata, "local_execution_runtime_miss_reason"), ), } } fn context_string(context: Option<&Map>, key: &str) -> Option { context .and_then(|value| value.get(key)) .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) } fn context_u64(context: Option<&Map>, key: &str) -> Option { context.and_then(|value| { value.get(key).and_then(|raw| { raw.as_u64() .or_else(|| raw.as_i64().and_then(|number| u64::try_from(number).ok())) }) }) } fn routing_u64_from_metadata(value: Option<&Value>, key: &str) -> Option { value .and_then(Value::as_object) .and_then(|object| object.get(key)) .and_then(|raw| { raw.as_u64() .or_else(|| raw.as_i64().and_then(|number| u64::try_from(number).ok())) }) } fn context_bool(context: Option<&Map>, key: &str) -> Option { context .and_then(|value| value.get(key)) .and_then(Value::as_bool) } fn context_value_ref<'a>(context: Option<&'a Map>, key: &str) -> Option<&'a Value> { context.and_then(|value| value.get(key)) } fn context_value(context: Option<&Map>, key: &str) -> Option { context.and_then(|value| value.get(key)).cloned() } fn context_usage_value(context: Option<&Map>, key: &str) -> Option { match context_value_ref(context, key) { Some(Value::Null) | None => None, Some(value) => clone_usage_capture_value(Some(value)), } } fn context_body_value(context: Option<&Map>, key: &str) -> Option { match context_value_ref(context, key) { Some(Value::Null) | None => None, Some(value) => clone_usage_body_value(Some(value)), } } fn context_has_inline_body(context: Option<&Map>, key: &str) -> bool { matches!(context_value_ref(context, key), Some(value) if !value.is_null()) } fn plan_has_inline_json_body_for_usage(plan: &ExecutionPlan) -> bool { plan.body.body_ref.is_none() && plan.body.body_bytes_b64.is_none() && plan.body.json_body.is_some() } fn build_runtime_request_capture_seed( plan: &ExecutionPlan, context: Option<&Map>, ) -> RuntimeRequestCaptureSeed { let request_body = context_body_value(context, "original_request_body"); let request_body_ref = context_string(context, "request_body_ref"); let provider_request = context_body_value(context, "provider_request_body") .or_else(|| plan_json_body_capture_for_usage(plan)); let provider_request_body_ref = context_string(context, "provider_request_body_ref") .or_else(|| non_empty_str(plan.body.body_ref.as_deref())); let body_states = build_runtime_body_states_seed_from_parts( request_body.is_some(), request_body_ref.as_deref(), provider_request.is_some(), provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.is_some(), ); RuntimeRequestCaptureSeed { request_body, request_body_ref, provider_request, provider_request_body_ref, body_states, } } fn build_runtime_request_metadata_seed( plan: &ExecutionPlan, context: Option<&Map>, ) -> Option { let request_body_ref = context_string(context, "request_body_ref"); let provider_request_body_ref = context_string(context, "provider_request_body_ref") .or_else(|| non_empty_str(plan.body.body_ref.as_deref())); let request_has_inline_body = context_has_inline_body(context, "original_request_body"); let provider_request_has_inline_body = context_has_inline_body(context, "provider_request_body") || plan_has_inline_json_body_for_usage(plan); let mut metadata = build_runtime_request_metadata_seed_from_parts( context, request_has_inline_body, request_body_ref.as_deref(), provider_request_has_inline_body, provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.as_deref(), ); if let Some(proxy) = plan.proxy.as_ref() { if let Some(node_id) = proxy .node_id .as_deref() .map(str::trim) .filter(|v| !v.is_empty()) { let mode = proxy.mode.as_deref().unwrap_or("").trim(); let mut proxy_obj = serde_json::Map::new(); proxy_obj.insert("node_id".to_string(), Value::String(node_id.to_string())); if !mode.is_empty() { proxy_obj.insert("mode".to_string(), Value::String(mode.to_string())); } let obj = metadata.get_or_insert_with(|| Value::Object(serde_json::Map::new())); if let Value::Object(map) = obj { map.insert("proxy".to_string(), Value::Object(proxy_obj)); } } } metadata } fn build_runtime_request_metadata_seed_from_parts( context: Option<&Map>, request_has_inline_body: bool, request_body_ref: Option<&str>, provider_request_has_inline_body: bool, provider_request_body_ref: Option<&str>, provider_request_body_base64: Option<&str>, ) -> Option { let mut metadata = Map::new(); if let Some(trace_id) = context_string(context, "trace_id") { metadata.insert("trace_id".to_string(), Value::String(trace_id)); } if let Some(client_ip) = context_string(context, "client_ip") { metadata.insert("client_ip".to_string(), Value::String(client_ip)); } if let Some(user_agent) = context_string(context, "user_agent") { metadata.insert("user_agent".to_string(), Value::String(user_agent)); } if let Some(client_requested_stream) = context_bool(context, "client_requested_stream") { metadata.insert( "client_requested_stream".to_string(), Value::Bool(client_requested_stream), ); } if let Some(upstream_is_stream) = context_bool(context, "upstream_is_stream") { metadata.insert( "upstream_is_stream".to_string(), Value::Bool(upstream_is_stream), ); } if let Some(api_key_is_standalone) = context_bool(context, "api_key_is_standalone") { metadata.insert( "api_key_is_standalone".to_string(), Value::Bool(api_key_is_standalone), ); } let provider_source_bytes = provider_request_body_base64.and_then(decoded_base64_len_hint); append_runtime_body_capture_metadata( &mut metadata, RuntimeBodyCaptureMetadataInput { request_has_inline_body, request_body_ref, provider_request_has_inline_body, provider_request_body_ref, provider_request_source_bytes: provider_source_bytes, provider_request_unavailable: provider_request_body_base64.is_some(), provider_request_unavailable_reason: provider_request_body_base64 .as_ref() .map(|_| "body_bytes_base64_only"), }, ); crate::body_capture::append_plan_body_capture_metadata( &mut metadata, provider_request_body_base64, ); (!metadata.is_empty()).then_some(Value::Object(metadata)) } fn capture_usage_storage_value(value: Value) -> Value { if usage_capture_within_limits(&value) { return value; } json!({ "truncated": true, "reason": "usage_capture_limits_exceeded", "max_depth": MAX_USAGE_CAPTURE_DEPTH, "max_nodes": MAX_USAGE_CAPTURE_NODES, "max_bytes": MAX_USAGE_CAPTURE_BYTES, "value_kind": usage_value_kind(&value), }) } fn build_runtime_body_states_seed( plan: &ExecutionPlan, context: Option<&Map>, ) -> UsageBodyStatesSeed { let request_body_ref = context_string(context, "request_body_ref"); let provider_request_body_ref = context_string(context, "provider_request_body_ref") .or_else(|| non_empty_str(plan.body.body_ref.as_deref())); build_runtime_body_states_seed_from_parts( context_has_inline_body(context, "original_request_body"), request_body_ref.as_deref(), context_has_inline_body(context, "provider_request_body") || plan_has_inline_json_body_for_usage(plan), provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.is_some(), ) } fn build_runtime_body_states_seed_from_parts( request_has_inline_body: bool, request_body_ref: Option<&str>, provider_request_has_inline_body: bool, provider_request_body_ref: Option<&str>, provider_request_unavailable: bool, ) -> UsageBodyStatesSeed { let states = build_runtime_body_capture_states( request_has_inline_body, request_body_ref, provider_request_has_inline_body, provider_request_body_ref, provider_request_unavailable, ); UsageBodyStatesSeed { request_body_state: Some(states.request), provider_request_body_state: Some(states.provider_request), response_body_state: Some(UsageBodyCaptureState::None), client_response_body_state: Some(UsageBodyCaptureState::None), } } fn merge_body_refs_seed_with_metadata( seed: &UsageBodyRefsSeed, metadata: Option<&Value>, ) -> UsageBodyRefsSeed { let object = metadata.and_then(Value::as_object); UsageBodyRefsSeed { request_body_ref: seed .request_body_ref .clone() .or_else(|| context_string(object, "request_body_ref")), provider_request_body_ref: seed .provider_request_body_ref .clone() .or_else(|| context_string(object, "provider_request_body_ref")), response_body_ref: seed .response_body_ref .clone() .or_else(|| context_string(object, "response_body_ref")), client_response_body_ref: seed .client_response_body_ref .clone() .or_else(|| context_string(object, "client_response_body_ref")), } } fn merge_body_refs_seed_with_metadata_owned( seed: UsageBodyRefsSeed, metadata: Option<&Value>, ) -> UsageBodyRefsSeed { let object = metadata.and_then(Value::as_object); UsageBodyRefsSeed { request_body_ref: seed .request_body_ref .or_else(|| context_string(object, "request_body_ref")), provider_request_body_ref: seed .provider_request_body_ref .or_else(|| context_string(object, "provider_request_body_ref")), response_body_ref: seed .response_body_ref .or_else(|| context_string(object, "response_body_ref")), client_response_body_ref: seed .client_response_body_ref .or_else(|| context_string(object, "client_response_body_ref")), } } fn plan_json_body_capture_for_usage(plan: &ExecutionPlan) -> Option { if plan.body.body_ref.is_some() || plan.body.body_bytes_b64.is_some() { return None; } clone_usage_body_value(plan.body.json_body.as_ref()) } fn clone_usage_capture_value(value: Option<&Value>) -> Option { value.cloned().map(capture_usage_storage_value) } fn clone_usage_body_value(value: Option<&Value>) -> Option { value.cloned() } fn sanitize_usage_event_capture_fields(mut data: UsageEventData) -> UsageEventData { data.request_headers = sanitize_usage_header_capture(data.request_headers); data.request_body_ref = sanitize_usage_body_ref(data.request_body_ref); data.provider_request_headers = sanitize_usage_header_capture(data.provider_request_headers); data.provider_request_body_ref = sanitize_usage_body_ref(data.provider_request_body_ref); data.response_headers = sanitize_usage_header_capture(data.response_headers); data.response_body_ref = sanitize_usage_body_ref(data.response_body_ref); data.client_response_headers = sanitize_usage_header_capture(data.client_response_headers); data.client_response_body_ref = sanitize_usage_body_ref(data.client_response_body_ref); data } fn sanitize_usage_event_capture_fields_trusted(mut data: UsageEventData) -> UsageEventData { data.request_headers = capture_usage_header_capture(data.request_headers); data.provider_request_headers = capture_usage_header_capture(data.provider_request_headers); data.response_headers = capture_usage_header_capture(data.response_headers); data.client_response_headers = capture_usage_header_capture(data.client_response_headers); data } fn sanitize_usage_event_data(mut data: UsageEventData) -> UsageEventData { data = sanitize_usage_event_capture_fields(data); data.request_metadata = sanitize_usage_request_metadata(data.request_metadata); data } fn capture_usage_header_capture(value: Option) -> Option { value.map(capture_usage_storage_value) } fn sanitize_usage_header_capture(value: Option) -> Option { mask_sensitive_headers_in_json_value(value).map(capture_usage_storage_value) } fn trim_owned_non_empty_string(value: String) -> Option { let trimmed = value.trim(); if trimmed.is_empty() { return None; } if trimmed.len() == value.len() { return Some(value); } Some(trimmed.to_string()) } fn sanitize_usage_body_ref(value: Option) -> Option { value.and_then(trim_owned_non_empty_string) } fn usage_capture_within_limits(value: &Value) -> bool { let mut nodes = 0usize; let mut estimated_bytes = 0usize; let mut stack = vec![(value, 1usize)]; while let Some((current, depth)) = stack.pop() { nodes = nodes.saturating_add(1); estimated_bytes = estimated_bytes.saturating_add(usage_value_size_hint(current)); if depth > MAX_USAGE_CAPTURE_DEPTH || nodes > MAX_USAGE_CAPTURE_NODES || estimated_bytes > MAX_USAGE_CAPTURE_BYTES { return false; } match current { Value::Array(items) => { estimated_bytes = estimated_bytes.saturating_add(items.len().saturating_mul(2)); for item in items.iter().rev() { stack.push((item, depth + 1)); } } Value::Object(object) => { estimated_bytes = estimated_bytes .saturating_add(object.len().saturating_mul(3)) .saturating_add( object .keys() .map(|key| key.len().saturating_add(2)) .sum::(), ); for item in object.values() { stack.push((item, depth + 1)); } } _ => {} } } true } fn usage_value_kind(value: &Value) -> &'static str { match value { Value::Null => "null", Value::Bool(_) => "bool", Value::Number(_) => "number", Value::String(_) => "string", Value::Array(_) => "array", Value::Object(_) => "object", } } fn usage_value_size_hint(value: &Value) -> usize { match value { Value::Null => 4, Value::Bool(false) => 5, Value::Bool(true) => 4, Value::Number(number) => number.to_string().len(), Value::String(text) => text.len().saturating_add(2), Value::Array(_) | Value::Object(_) => 2, } } fn non_empty_str(value: Option<&str>) -> Option { value .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) } fn infer_request_type(api_format: Option<&str>) -> String { match infer_endpoint_kind(api_format.unwrap_or_default()) { Some("video") => "video".to_string(), Some("image") => "image".to_string(), _ => "chat".to_string(), } } fn infer_api_family(api_format: &str) -> Option<&str> { api_format.split_once(':').map(|(family, _)| family) } fn infer_endpoint_kind(api_format: &str) -> Option<&str> { api_format.split_once(':').map(|(_, kind)| kind) } fn apply_standardized_usage_seed(usage: &StandardizedUsage, data: &mut UsageEventData) { if usage.input_tokens > 0 { data.input_tokens = Some(usage.input_tokens as u64); } if usage.output_tokens > 0 { data.output_tokens = Some(usage.output_tokens as u64); } if usage.cache_creation_tokens > 0 { data.cache_creation_input_tokens = Some(usage.cache_creation_tokens as u64); } if usage.cache_creation_ephemeral_5m_tokens > 0 { data.cache_creation_ephemeral_5m_input_tokens = Some(usage.cache_creation_ephemeral_5m_tokens as u64); } if usage.cache_creation_ephemeral_1h_tokens > 0 { data.cache_creation_ephemeral_1h_input_tokens = Some(usage.cache_creation_ephemeral_1h_tokens as u64); } if usage.cache_read_tokens > 0 { data.cache_read_input_tokens = Some(usage.cache_read_tokens as u64); } let total_tokens = standardized_usage_total_tokens(usage); if total_tokens > 0 { data.total_tokens = Some(total_tokens); } } fn standardized_usage_total_tokens(usage: &StandardizedUsage) -> u64 { positive_usage_component(usage.input_tokens) .saturating_add(positive_usage_component(usage.output_tokens)) .saturating_add(positive_usage_component(usage.cache_creation_tokens)) .saturating_add(positive_usage_component(usage.cache_read_tokens)) } fn positive_usage_component(value: i64) -> u64 { value.max(0) as u64 } fn headers_to_json(headers: &BTreeMap) -> Option { if headers.is_empty() { return None; } Some(Value::Object(Map::from_iter(headers.iter().map( |(key, value)| (key.clone(), Value::String(mask_header_value(key, value))), )))) } /// 默认敏感请求头清单。与 /// `apps/aether-gateway/src/handlers/admin/system/shared/configs.rs` 中 /// `sensitive_headers` 系统配置默认值保持一致。 const DEFAULT_SENSITIVE_HEADERS: &[&str] = &[ "authorization", "x-api-key", "api-key", "x-goog-api-key", "cookie", "set-cookie", "proxy-authorization", ]; /// 判断 header 名是否属于敏感字段（大小写不敏感）。 fn is_sensitive_header(name: &str) -> bool { let trimmed = name.trim(); DEFAULT_SENSITIVE_HEADERS .iter() .any(|candidate| trimmed.eq_ignore_ascii_case(candidate)) } /// 对单个 header value 进行脱敏：保留前 4 + 后 4 字符，中间替换为 `****`。 /// 长度小于等于 8 时整体替换为 `****`。 fn mask_header_value(name: &str, value: &str) -> String { if !is_sensitive_header(name) { return value.to_string(); } mask_sensitive_header_value(value) } fn mask_sensitive_header_value(value: &str) -> String { if value.len() <= 8 { return "****".to_string(); } let prefix: String = value.chars().take(4).collect(); let suffix: String = value .chars() .rev() .take(4) .collect::>() .into_iter() .rev() .collect(); format!("{prefix}****{suffix}") } /// 对 JSON 形式的 headers 做就地脱敏。仅当 value 是 Object 时才会处理； /// 其它形式的值保持不变。 fn mask_sensitive_headers_in_json_value(value: Option) -> Option { let mut value = value?; let Value::Object(map) = &mut value else { return Some(value); }; for (key, val) in map.iter_mut() { if !is_sensitive_header(key) { continue; } match val { Value::String(text) => { *text = mask_sensitive_header_value(text); } Value::Null => {} other => { *other = Value::String(mask_sensitive_header_value(&other.to_string())); } } } Some(value) } fn resolve_error_category(status_code: u16, event_type: UsageEventType) -> Option { match event_type { UsageEventType::Cancelled => Some("cancelled".to_string()), UsageEventType::Failed if status_code >= 500 => Some("server_error".to_string()), UsageEventType::Failed if status_code >= 400 => Some("client_error".to_string()), UsageEventType::Failed if status_code >= 300 => Some("redirect".to_string()), UsageEventType::Failed => Some("non_success_status".to_string()), _ => None, } } fn resolve_error_message( status_code: u16, body_json: Option<&Value>, body_base64: Option<&str>, ) -> Option { let decoded_body = body_json .is_none() .then(|| decode_body_for_storage(body_base64)) .flatten(); let error_body = body_json.or(decoded_body.as_ref()); let explicit_error_message = error_body.and_then(extract_explicit_error_message_from_json); if explicit_error_message.is_some() { return explicit_error_message; } if (200..300).contains(&status_code) { return None; } error_body.and_then(extract_generic_error_message_from_json) } fn extract_explicit_error_message_from_json(value: &Value) -> Option { value .get("error") .and_then(|error| error.get("message")) .and_then(Value::as_str) .map(ToOwned::to_owned) } fn extract_generic_error_message_from_json(value: &Value) -> Option { extract_explicit_error_message_from_json(value).or_else(|| { value .get("message") .and_then(Value::as_str) .map(ToOwned::to_owned) }) } fn decode_body_for_storage(body_base64: Option<&str>) -> Option { let body_base64 = body_base64?; let bytes = base64::engine::general_purpose::STANDARD .decode(body_base64) .ok()?; if let Ok(json_body) = serde_json::from_slice::(&bytes) { return Some(json_body); } if let Ok(text) = String::from_utf8(bytes) { if let Some(stream_body) = parse_sse_body_for_storage(&text) { return Some(stream_body); } return Some(Value::String(text)); } Some(Value::String(body_base64.to_string())) } fn flush_sse_payload(payload: &mut String, has_payload: &mut bool, on_payload: &mut F) where F: FnMut(&str), { if !*has_payload { return; } on_payload(payload); payload.clear(); *has_payload = false; } fn for_each_sse_payload(text: &str, mut on_payload: F) where F: FnMut(&str), { let mut payload = String::new(); let mut has_payload = false; let bytes = text.as_bytes(); let mut line_start = 0usize; let mut cursor = 0usize; while cursor <= bytes.len() { if cursor < bytes.len() && bytes[cursor] != b'\n' && bytes[cursor] != b'\r' { cursor += 1; continue; } let line = text[line_start..cursor].trim(); if line.is_empty() { flush_sse_payload(&mut payload, &mut has_payload, &mut on_payload); } else if let Some(data) = line.strip_prefix("data:").map(str::trim) { if !data.is_empty() { if has_payload { payload.push('\n'); } payload.push_str(data); has_payload = true; } } if cursor == bytes.len() { break; } if bytes[cursor] == b'\r' && bytes.get(cursor + 1) == Some(&b'\n') { cursor += 2; } else { cursor += 1; } line_start = cursor; } flush_sse_payload(&mut payload, &mut has_payload, &mut on_payload); } fn parse_sse_body_for_storage(text: &str) -> Option { let mut chunks = Vec::new(); let mut total_chunks = 0_u64; let mut saw_done = false; for_each_sse_payload(text, |payload| { if payload == "[DONE]" { saw_done = true; return; } total_chunks += 1; if let Ok(json_body) = serde_json::from_str::(payload) { chunks.push(json_body); } }); if total_chunks == 0 && !saw_done { return None; } let stored_chunks = chunks.len() as u64; let mut metadata = Map::from_iter([ ("stream".to_string(), Value::Bool(true)), ("total_chunks".to_string(), json!(total_chunks)), ("stored_chunks".to_string(), json!(stored_chunks)), ("content_length".to_string(), json!(text.len())), ]); if saw_done { metadata.insert("has_completion".to_string(), Value::Bool(true)); } if stored_chunks < total_chunks { metadata.insert( "dropped_chunks".to_string(), json!(total_chunks - stored_chunks), ); } if chunks.is_empty() { metadata.insert( "parse_error".to_string(), Value::String("Failed to parse response as SSE JSON format".to_string()), ); return Some(json!({ "chunks": [], "raw_response": text, "metadata": metadata, })); } Some(json!({ "chunks": chunks, "metadata": metadata, })) } fn extract_token_counts_from_sse_text(text: &str) -> Option<(u64, u64, u64)> { let mut last_seen = None; for_each_sse_payload(text, |payload| { if payload == "[DONE]" { return; } if let Ok(json_body) = serde_json::from_str::(payload) { if let Some(tokens) = extract_token_counts_from_json(&json_body) { last_seen = Some(tokens); } } }); last_seen } fn extract_token_counts_from_value(value: &Value) -> Option<(u64, u64, u64)> { match value { Value::String(text) => extract_token_counts_from_sse_text(text), _ => extract_token_counts_from_json(value), } } fn extract_token_counts_from_json(value: &Value) -> Option<(u64, u64, u64)> { if let Some(usage) = value.get("usage").and_then(Value::as_object) { let usage_details = usage .get("input_tokens_details") .or_else(|| usage.get("prompt_tokens_details")) .and_then(Value::as_object); let input = usage .get("input_tokens") .or_else(|| usage.get("prompt_tokens")) .and_then(Value::as_u64) .unwrap_or_default(); let output = usage .get("output_tokens") .or_else(|| usage.get("completion_tokens")) .and_then(Value::as_u64) .unwrap_or_default(); let raw_total = usage .get("total_tokens") .and_then(Value::as_u64) .unwrap_or(input + output); let cache_creation = extract_cache_creation_tokens_from_usage_object(usage, usage_details); let cache_read = extract_cache_read_tokens_from_usage_object(usage, usage_details); let total = if cache_creation > 0 || cache_read > 0 { input .saturating_add(output) .saturating_add(cache_creation) .saturating_add(cache_read) } else { raw_total }; return Some((input, output, total)); } if let Some(usage) = value.get("usageMetadata").and_then(Value::as_object) { let input = usage .get("promptTokenCount") .and_then(Value::as_u64) .unwrap_or_default(); let output = usage .get("candidatesTokenCount") .and_then(Value::as_u64) .unwrap_or_default(); let raw_total = usage .get("totalTokenCount") .and_then(Value::as_u64) .unwrap_or(input + output); return Some((input, output, raw_total)); } if let Some(chunks) = value.get("chunks").and_then(Value::as_array) { if let Some(tokens) = chunks.iter().rev().find_map(extract_token_counts_from_json) { return Some(tokens); } } if let Some(response) = value.get("response") { return extract_token_counts_from_json(response); } None } fn extract_cache_creation_tokens_from_usage_object( usage: &serde_json::Map, usage_details: Option<&serde_json::Map>, ) -> u64 { let direct = usage .get("cache_creation_input_tokens") .and_then(Value::as_u64) .unwrap_or_default(); if direct > 0 { return direct; } let breakdown = usage .get("cache_creation") .and_then(Value::as_object) .map(|cache_creation| { cache_creation .get("ephemeral_5m_input_tokens") .and_then(Value::as_u64) .unwrap_or_default() .saturating_add( cache_creation .get("ephemeral_1h_input_tokens") .and_then(Value::as_u64) .unwrap_or_default(), ) }) .unwrap_or_default(); if breakdown > 0 { return breakdown; } usage_details .and_then(|details| details.get("cached_creation_tokens")) .and_then(Value::as_u64) .unwrap_or_default() } fn extract_cache_read_tokens_from_usage_object( usage: &serde_json::Map, usage_details: Option<&serde_json::Map>, ) -> u64 { usage .get("cache_read_input_tokens") .and_then(Value::as_u64) .or_else(|| { usage_details .and_then(|details| details.get("cached_tokens")) .and_then(Value::as_u64) }) .unwrap_or_default() } fn empty_to_none(value: Option) -> Option { value.and_then(trim_owned_non_empty_string) } #[cfg(test)] mod tests { use super::{ build_pending_usage_record, build_pending_usage_record_from_seed, build_stream_terminal_usage_event, build_streaming_usage_record, build_sync_terminal_usage_event, build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed, build_terminal_usage_context_seed, build_terminal_usage_event_from_seed, build_usage_event_data_seed, extract_token_counts_from_json, extract_token_counts_from_value, headers_to_json, mask_header_value, mask_sensitive_headers_in_json_value, parse_sse_body_for_storage, resolve_error_message, trim_owned_non_empty_string, LifecycleUsageSeed, TerminalUsageSeed, UsageBodyRefsSeed, UsageBodyStatesSeed, UsageRoutingSeed, UsageTerminalState, MAX_USAGE_CAPTURE_BYTES, MAX_USAGE_CAPTURE_DEPTH, }; use crate::{ build_upsert_usage_record_from_event, GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEvent, UsageEventData, UsageEventType, }; use aether_contracts::{ ExecutionPlan, ExecutionStreamTerminalSummary, RequestBody, StandardizedUsage, }; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use base64::Engine as _; use serde_json::{json, Value}; use std::collections::BTreeMap; #[test] fn extracts_openai_usage_tokens() { let tokens = extract_token_counts_from_json(&json!({ "usage": { "input_tokens": 3, "output_tokens": 5, "total_tokens": 8 } })) .expect("tokens should exist"); assert_eq!(tokens, (3, 5, 8)); } #[test] fn extracts_openai_usage_tokens_with_cache_components() { let tokens = extract_token_counts_from_json(&json!({ "usage": { "input_tokens": 3, "output_tokens": 5, "total_tokens": 8, "input_tokens_details": { "cached_tokens": 2, "cached_creation_tokens": 1 } } })) .expect("tokens should exist"); assert_eq!(tokens, (3, 5, 11)); } #[test] fn extracts_openai_usage_tokens_with_prompt_token_details() { let tokens = extract_token_counts_from_json(&json!({ "usage": { "prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8, "prompt_tokens_details": { "cached_tokens": 2, "cached_creation_tokens": 1 } } })) .expect("tokens should exist"); assert_eq!(tokens, (3, 5, 11)); } #[test] fn extracts_claude_usage_tokens_with_cache_components() { let tokens = extract_token_counts_from_json(&json!({ "usage": { "input_tokens": 6, "output_tokens": 20, "cache_creation_input_tokens": 41857, "cache_read_input_tokens": 0 } })) .expect("tokens should exist"); assert_eq!(tokens, (6, 20, 41883)); } #[test] fn extracts_gemini_usage_tokens_without_adding_cached_content_twice() { let tokens = extract_token_counts_from_json(&json!({ "usageMetadata": { "promptTokenCount": 14, "candidatesTokenCount": 6, "cachedContentTokenCount": 2, "totalTokenCount": 20 } })) .expect("tokens should exist"); assert_eq!(tokens, (14, 6, 20)); } #[test] fn extracts_usage_tokens_from_last_matching_chunk() { let tokens = extract_token_counts_from_json(&json!({ "chunks": [ { "response": { "usage": { "input_tokens": 1, "output_tokens": 2, "total_tokens": 3 } } }, { "response": { "usage": { "input_tokens": 9, "output_tokens": 4, "total_tokens": 13, "input_tokens_details": { "cached_tokens": 5, "cached_creation_tokens": 2 } } } } ] })) .expect("tokens should exist"); assert_eq!(tokens, (9, 4, 20)); } #[test] fn trim_owned_non_empty_string_preserves_clean_values_and_drops_blank_ones() { assert_eq!( trim_owned_non_empty_string("body-ref-1".to_string()), Some("body-ref-1".to_string()), ); assert_eq!( trim_owned_non_empty_string(" body-ref-1 ".to_string()), Some("body-ref-1".to_string()), ); assert_eq!(trim_owned_non_empty_string(" ".to_string()), None); } #[test] fn builds_upsert_record_from_terminal_event() { let record = build_upsert_usage_record_from_event(&UsageEvent { event_type: UsageEventType::Completed, request_id: "req-1".to_string(), timestamp_ms: 1_700_000_000_000, data: UsageEventData { user_id: Some("user-1".to_string()), api_key_id: Some("key-1".to_string()), provider_name: "OpenAI".to_string(), model: "gpt-5".to_string(), api_format: Some("openai:chat".to_string()), endpoint_api_format: Some("openai:chat".to_string()), input_tokens: Some(10), output_tokens: Some(20), total_tokens: Some(30), status_code: Some(200), ..UsageEventData::default() }, }) .expect("record should build"); assert_eq!(record.request_id, "req-1"); assert_eq!(record.status, "completed"); assert_eq!(record.billing_status, "pending"); assert_eq!(record.total_tokens, Some(30)); } #[test] fn pending_usage_records_stay_lightweight() { let plan = ExecutionPlan { request_id: "req-pending-usage-1".to_string(), candidate_id: Some("cand-pending-usage-1".to_string()), provider_name: Some("Codex Proxy".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/messages".to_string(), headers: BTreeMap::from([( "authorization".to_string(), "Bearer pending-secret".to_string(), )]), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}] })), stream: false, client_api_format: "claude:messages".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let record = build_pending_usage_record( &plan, Some(&json!({ "candidate_id": "cand-pending-usage-1", "candidate_index": 0, "key_name": "codex-upstream", "planner_kind": "claude_cli_sync", "route_family": "claude", "route_kind": "cli", "execution_path": "local_execution_runtime_miss", "local_execution_runtime_miss_reason": "all_candidates_skipped", "original_headers": {"authorization": "Bearer upstream-secret"}, "original_request_body": {"messages": [{"content": "should not be persisted in pending"}]}, "provider_request_headers": {"authorization": "Bearer provider-secret"}, "provider_request_body": {"input": "should not be persisted in pending"} })), 1_700_000_000, ) .expect("pending usage should build"); assert_eq!(record.status, "pending"); assert_eq!(record.billing_status, "pending"); assert!(record.request_headers.is_none()); assert!(record.request_body.is_none()); assert!(record.provider_request_headers.is_none()); assert!(record.provider_request_body.is_none()); assert_eq!(record.candidate_id.as_deref(), Some("cand-pending-usage-1")); assert_eq!(record.candidate_index, Some(0)); assert_eq!(record.key_name.as_deref(), Some("codex-upstream")); assert_eq!(record.planner_kind.as_deref(), Some("claude_cli_sync")); assert_eq!(record.route_family.as_deref(), Some("claude")); assert_eq!(record.route_kind.as_deref(), Some("cli")); assert_eq!( record.execution_path.as_deref(), Some("local_execution_runtime_miss") ); assert_eq!( record.local_execution_runtime_miss_reason.as_deref(), Some("all_candidates_skipped") ); assert_eq!(record.request_metadata, None); } #[test] fn pending_usage_record_preserves_standalone_key_metadata() { let plan = ExecutionPlan { request_id: "req-pending-standalone-1".to_string(), candidate_id: None, provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"model": "gpt-5.4"})), stream: false, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let record = build_pending_usage_record( &plan, Some(&json!({ "api_key_is_standalone": true, "client_ip": "203.0.113.8", "user_agent": "Claude-Code/1.0" })), 1_700_000_000, ) .expect("pending usage should build"); assert_eq!( record.request_metadata, Some(json!({ "api_key_is_standalone": true, "client_ip": "203.0.113.8", "user_agent": "Claude-Code/1.0" })) ); } #[test] fn streaming_usage_records_stay_lightweight_by_default() { let plan = ExecutionPlan { request_id: "req-streaming-usage-1".to_string(), candidate_id: Some("cand-streaming-usage-1".to_string()), provider_name: Some("Codex Proxy".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/messages".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"model": "gpt-5.4"})), stream: true, client_api_format: "claude:messages".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let record = build_streaming_usage_record( &plan, Some(&json!({ "candidate_id": "cand-streaming-usage-1", "original_request_body": {"messages": [{"content": "omit me"}]}, "provider_request_body": {"input": "omit me too"} })), 200, None, 1_700_000_010, ) .expect("streaming usage should build"); assert_eq!(record.status, "streaming"); assert_eq!(record.billing_status, "pending"); assert_eq!(record.status_code, Some(200)); assert!(record.request_headers.is_none()); assert!(record.request_body.is_none()); assert!(record.provider_request_headers.is_none()); assert!(record.provider_request_body.is_none()); assert!(record.response_headers.is_none()); assert!(record.client_response_headers.is_none()); } #[test] fn sync_terminal_usage_preserves_overdeep_request_payloads_for_repo_body_storage() { const DEEP_NESTED_LEVELS: usize = MAX_USAGE_CAPTURE_DEPTH + 8; let mut nested = Value::String("leaf".to_string()); for depth in 0..DEEP_NESTED_LEVELS { nested = json!({ "depth": depth, "child": nested }); } let plan = ExecutionPlan { request_id: "req-sync-usage-deep-1".to_string(), candidate_id: Some("cand-sync-usage-deep-1".to_string()), provider_name: Some("Codex Proxy".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/messages".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"model": "gpt-5.4"})), stream: false, client_api_format: "claude:messages".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-usage-deep-1".to_string(), report_kind: "claude_cli_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "claude:messages", "provider_api_format": "openai:responses", "needs_conversion": true, "original_request_body": nested, "provider_request_body": {"input": "safe"} })), status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({"id": "resp-1"})), body_base64: None, client_body_json: Some(json!({"type": "message"})), telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("terminal usage should build"); assert_eq!( event.data.request_body, payload .report_context .as_ref() .and_then(|value| value.get("original_request_body")) .cloned() ); assert_eq!( event.data.provider_request_body, Some(json!({"input": "safe"})) ); } #[test] fn builds_stream_terminal_usage_from_provider_body_and_preserves_client_body() { let plan = ExecutionPlan { request_id: "req-stream-usage-1".to_string(), candidate_id: Some("cand-stream-usage-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:chat".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewayStreamReportRequest { trace_id: "trace-stream-usage-1".to_string(), report_kind: "openai_chat_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "openai:responses", "needs_conversion": true })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: Some( base64::engine::general_purpose::STANDARD.encode( serde_json::to_vec(&json!({ "usage": { "prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8 } })) .expect("provider body should encode"), ), ), provider_body_state: Some(UsageBodyCaptureState::Inline), client_body_base64: Some( base64::engine::general_purpose::STANDARD .encode("data: {\"id\":\"chatcmpl_123\"}\n\ndata: [DONE]\n"), ), client_body_state: Some(UsageBodyCaptureState::Inline), terminal_summary: None, telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(3)); assert_eq!(event.data.output_tokens, Some(5)); assert_eq!(event.data.total_tokens, Some(8)); assert_eq!( event.data.response_body, Some(json!({ "usage": { "prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8 } })) ); assert_eq!( event.data.client_response_body, Some(json!({ "chunks": [ { "id": "chatcmpl_123" } ], "metadata": { "stream": true, "total_chunks": 1, "stored_chunks": 1, "content_length": 42, "has_completion": true } })) ); } #[test] fn stream_terminal_usage_marks_redirect_status_as_failed() { let plan = ExecutionPlan { request_id: "req-stream-redirect-usage".to_string(), candidate_id: Some("cand-stream-redirect-usage".to_string()), provider_name: Some("ChatGPTWeb".to_string()), provider_id: "provider-redirect".to_string(), endpoint_id: "endpoint-redirect".to_string(), key_id: "key-redirect".to_string(), method: "POST".to_string(), url: "https://example.com/v1beta/models/gemini:streamGenerateContent".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "gemini:generate_content".to_string(), provider_api_format: "gemini:generate_content".to_string(), model_name: Some("gemini".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let client_body = json!({ "error": { "type": "execution_runtime_non_success_status", "message": "execution runtime stream returned non-success status 302", "code": 302, "upstream_status": 302, "location": "/" } }); let payload = GatewayStreamReportRequest { trace_id: "trace-stream-redirect-usage".to_string(), report_kind: "gemini_chat_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "gemini:generate_content", "provider_api_format": "gemini:generate_content" })), status_code: 302, headers: BTreeMap::from([ ("content-type".to_string(), "application/json".to_string()), ("x-aether-upstream-status".to_string(), "302".to_string()), ]), provider_body_base64: Some( base64::engine::general_purpose::STANDARD .encode(br#"{"error":{"message":"raw redirect body"}}"#), ), provider_body_state: Some(UsageBodyCaptureState::Inline), client_body_base64: Some( base64::engine::general_purpose::STANDARD .encode(serde_json::to_vec(&client_body).expect("body should encode")), ), client_body_state: Some(UsageBodyCaptureState::Inline), terminal_summary: None, telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.event_type, UsageEventType::Failed); assert_eq!(event.data.status_code, Some(302)); assert_eq!(event.data.error_category.as_deref(), Some("redirect")); assert_eq!( event.data.error_message.as_deref(), Some("raw redirect body") ); } #[test] fn builds_stream_terminal_usage_from_terminal_summary_usage_without_decoding_bodies() { let plan = ExecutionPlan { request_id: "req-stream-summary-usage-1".to_string(), candidate_id: Some("cand-stream-summary-usage-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:chat".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let mut standardized_usage = StandardizedUsage::new(); standardized_usage.input_tokens = 13; standardized_usage.output_tokens = 21; standardized_usage.cache_creation_tokens = 2; standardized_usage.cache_read_tokens = 3; let payload = GatewayStreamReportRequest { trace_id: "trace-stream-summary-usage-1".to_string(), report_kind: "openai_chat_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "openai:responses", "needs_conversion": true })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: None, provider_body_state: Some(UsageBodyCaptureState::None), client_body_base64: None, client_body_state: Some(UsageBodyCaptureState::None), terminal_summary: Some(ExecutionStreamTerminalSummary { standardized_usage: Some(standardized_usage), finish_reason: Some("stop".to_string()), response_id: Some("resp_summary_1".to_string()), model: Some("gpt-5.4".to_string()), observed_finish: true, unknown_event_count: 0, parser_error: None, }), telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(13)); assert_eq!(event.data.output_tokens, Some(21)); assert_eq!(event.data.total_tokens, Some(39)); assert_eq!(event.data.cache_creation_input_tokens, Some(2)); assert_eq!(event.data.cache_read_input_tokens, Some(3)); assert!(event.data.response_body.is_none()); assert!(event.data.client_response_body.is_none()); } #[test] fn stream_terminal_usage_prefers_more_complete_provider_chunks_usage() { let plan = ExecutionPlan { request_id: "req-stream-provider-chunks-usage-1".to_string(), candidate_id: Some("cand-stream-provider-chunks-usage-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.5".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let mut partial_summary_usage = StandardizedUsage::new(); partial_summary_usage.output_tokens = 148; let provider_body = json!({ "chunks": [ { "type": "response.created", "response": { "id": "resp_123", "object": "response", "model": "gpt-5.5", "status": "in_progress", "usage": null } }, { "type": "response.completed", "response": { "id": "resp_123", "object": "response", "model": "gpt-5.5", "status": "completed", "usage": { "input_tokens": 26, "input_tokens_details": { "cached_tokens": 0 }, "output_tokens": 148, "output_tokens_details": { "reasoning_tokens": 10 }, "total_tokens": 174 } }, "sequence_number": 141 } ], "metadata": { "stream": true, "total_chunks": 142, "stored_chunks": 142 } }); let payload = GatewayStreamReportRequest { trace_id: "trace-stream-provider-chunks-usage-1".to_string(), report_kind: "openai_responses_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:responses", "provider_api_format": "openai:responses", })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: Some( base64::engine::general_purpose::STANDARD.encode(provider_body.to_string()), ), provider_body_state: Some(UsageBodyCaptureState::Inline), client_body_base64: None, client_body_state: Some(UsageBodyCaptureState::None), terminal_summary: Some(ExecutionStreamTerminalSummary { standardized_usage: Some(partial_summary_usage), finish_reason: None, response_id: Some("resp_123".to_string()), model: Some("gpt-5.5".to_string()), observed_finish: true, unknown_event_count: 0, parser_error: None, }), telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(26)); assert_eq!(event.data.output_tokens, Some(148)); assert_eq!(event.data.total_tokens, Some(174)); assert_eq!(event.data.cache_read_input_tokens, None); } #[test] fn builds_stream_terminal_usage_from_sse_chunks_and_extracts_usage() { let plan = ExecutionPlan { request_id: "req-stream-usage-2".to_string(), candidate_id: Some("cand-stream-usage-2".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let sse_body = concat!( "event: response.created\n", "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n", "event: response.output_text.delta\n", "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello from Responses stream\"}\n\n", "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello from Responses stream\"}]}],\"usage\":{\"input_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":2,\"cached_creation_tokens\":1},\"output_tokens\":5,\"output_tokens_details\":{\"reasoning_tokens\":1},\"total_tokens\":8}}}\n\n", "data: [DONE]\n", ); let payload = GatewayStreamReportRequest { trace_id: "trace-stream-usage-2".to_string(), report_kind: "openai_responses_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:responses", "provider_api_format": "openai:responses", })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: Some(base64::engine::general_purpose::STANDARD.encode(sse_body)), provider_body_state: Some(UsageBodyCaptureState::Inline), client_body_base64: None, client_body_state: Some(UsageBodyCaptureState::None), terminal_summary: None, telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(3)); assert_eq!(event.data.output_tokens, Some(5)); assert_eq!(event.data.total_tokens, Some(11)); assert_eq!(event.data.cache_creation_input_tokens, Some(1)); assert_eq!(event.data.cache_read_input_tokens, Some(2)); assert_eq!( event.data.response_body, Some(json!({ "chunks": [ { "type": "response.created", "response": { "id": "resp_123", "object": "response", "model": "gpt-5.4", "status": "in_progress" } }, { "type": "response.output_text.delta", "delta": "Hello from Responses stream" }, { "type": "response.completed", "response": { "id": "resp_123", "object": "response", "model": "gpt-5.4", "status": "completed", "output": [ { "type": "message", "role": "assistant", "content": [ { "type": "output_text", "text": "Hello from Responses stream" } ] } ], "usage": { "input_tokens": 3, "input_tokens_details": { "cached_tokens": 2, "cached_creation_tokens": 1 }, "output_tokens": 5, "output_tokens_details": { "reasoning_tokens": 1 }, "total_tokens": 8 } } } ], "metadata": { "stream": true, "total_chunks": 3, "stored_chunks": 3, "content_length": sse_body.len(), "has_completion": true } })) ); } #[test] fn builds_sync_terminal_usage_from_provider_body_and_preserves_client_body() { let plan = ExecutionPlan { request_id: "req-sync-usage-1".to_string(), candidate_id: Some("cand-sync-usage-1".to_string()), provider_name: Some("Gemini".to_string()), provider_id: "provider-2".to_string(), endpoint_id: "endpoint-2".to_string(), key_id: "key-2".to_string(), method: "POST".to_string(), url: "https://example.com/v1beta/models/gemini:generateContent".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: false, client_api_format: "openai:chat".to_string(), provider_api_format: "gemini:generate_content".to_string(), model_name: Some("gpt-5".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-usage-1".to_string(), report_kind: "openai_chat_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "gemini:generate_content", "needs_conversion": true })), status_code: 200, headers: BTreeMap::from([ ( "authorization".to_string(), "Bearer very-secret-token".to_string(), ), ("content-type".to_string(), "application/json".to_string()), ]), body_json: Some(json!({ "usageMetadata": { "promptTokenCount": 4, "candidatesTokenCount": 6, "totalTokenCount": 10 } })), client_body_json: Some(json!({ "id": "chatcmpl_456", "object": "chat.completion" })), body_base64: None, telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(4)); assert_eq!(event.data.output_tokens, Some(6)); assert_eq!(event.data.total_tokens, Some(10)); assert_eq!( event.data.response_headers, Some(json!({ "authorization": "Bear****oken", "content-type": "application/json" })) ); assert_eq!( event.data.client_response_headers, Some(json!({ "authorization": "Bear****oken", "content-type": "application/json" })) ); assert_eq!( event.data.response_body, Some(json!({ "usageMetadata": { "promptTokenCount": 4, "candidatesTokenCount": 6, "totalTokenCount": 10 } })) ); assert_eq!( event.data.client_response_body, Some(json!({ "id": "chatcmpl_456", "object": "chat.completion" })) ); } #[test] fn sync_terminal_usage_prefers_upstream_stream_body_over_aggregated_sync_body() { let sse_body = concat!( "event: response.created\n", "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_sync_stream_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\",\"output\":[]}}\n\n", "event: response.output_text.delta\n", "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello from upstream stream\"}\n\n", "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_sync_stream_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"total_tokens\":8}}}\n\n", "data: [DONE]\n", ); let plan = ExecutionPlan { request_id: "req-sync-upstream-stream-1".to_string(), candidate_id: Some("cand-sync-upstream-stream-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: false, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-upstream-stream-1".to_string(), report_kind: "openai_responses_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:responses", "provider_api_format": "openai:responses", "upstream_is_stream": true })), status_code: 200, headers: BTreeMap::from([( "content-type".to_string(), "text/event-stream".to_string(), )]), body_json: Some(json!({ "id": "resp_sync_stream_123", "object": "response", "status": "completed", "output": [], "usage": { "input_tokens": 1, "output_tokens": 1, "total_tokens": 2 } })), client_body_json: Some(json!({ "id": "resp_sync_stream_123", "object": "response", "status": "completed", "output": [{ "type": "message", "role": "assistant", "content": [{ "type": "output_text", "text": "Hello from upstream stream" }] }], "usage": { "input_tokens": 3, "output_tokens": 5, "total_tokens": 8 } })), body_base64: Some(base64::engine::general_purpose::STANDARD.encode(sse_body)), telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(3)); assert_eq!(event.data.output_tokens, Some(5)); assert_eq!(event.data.total_tokens, Some(8)); assert_eq!( event .data .response_body .as_ref() .and_then(|value| value.get("chunks")) .and_then(Value::as_array) .and_then(|chunks| chunks.get(1)) .and_then(|chunk| chunk.get("type")) .and_then(Value::as_str), Some("response.output_text.delta") ); assert_eq!( event.data.client_response_body.as_ref().and_then(|value| { value .get("output") .and_then(Value::as_array) .and_then(|output| output.first()) .and_then(|item| item.get("content")) .and_then(Value::as_array) .and_then(|content| content.first()) .and_then(|part| part.get("text")) .and_then(Value::as_str) }), Some("Hello from upstream stream") ); } #[test] fn sync_terminal_seed_path_matches_legacy_wrapper_event() { let plan = ExecutionPlan { request_id: "req-sync-seed-match-1".to_string(), candidate_id: Some("cand-sync-seed-match-1".to_string()), provider_name: Some("Gemini".to_string()), provider_id: "provider-2".to_string(), endpoint_id: "endpoint-2".to_string(), key_id: "key-2".to_string(), method: "POST".to_string(), url: "https://example.com/v1beta/models/gemini:generateContent".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: false, client_api_format: "openai:chat".to_string(), provider_api_format: "gemini:generate_content".to_string(), model_name: Some("gpt-5".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-seed-match-1".to_string(), report_kind: "openai_chat_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "gemini:generate_content", "needs_conversion": true })), status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({ "usageMetadata": { "promptTokenCount": 4, "candidatesTokenCount": 6, "totalTokenCount": 10 } })), client_body_json: Some(json!({ "id": "chatcmpl_456", "object": "chat.completion" })), body_base64: None, telemetry: None, }; let legacy_event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("legacy terminal usage should build"); let context_seed = build_terminal_usage_context_seed(&plan, payload.report_context.as_ref()); let payload_seed = build_sync_terminal_usage_payload_seed(&payload); let seed_event = build_terminal_usage_event_from_seed(build_sync_terminal_usage_seed( context_seed, payload_seed, )) .expect("seed terminal usage should build"); assert_eq!(seed_event.event_type, legacy_event.event_type); assert_eq!(seed_event.request_id, legacy_event.request_id); assert_eq!(seed_event.data, legacy_event.data); } #[test] fn sync_terminal_usage_exposes_provider_request_body_ref_as_typed_field() { let plan = ExecutionPlan { request_id: "req-sync-body-ref-1".to_string(), candidate_id: Some("cand-sync-body-ref-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody { json_body: Some(json!({ "model": "gpt-5.4", "input": "this should stay out of usage storage" })), body_bytes_b64: None, body_ref: Some("blob://provider-request-1".to_string()), }, stream: false, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-body-ref-1".to_string(), report_kind: "openai_responses_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:responses", "provider_api_format": "openai:responses", "trace_id": "trace-sync-body-ref-1" })), status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({"id": "resp_123"})), client_body_json: None, body_base64: None, telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert!(event.data.request_body.is_none()); assert!(event.data.provider_request_body.is_none()); assert_eq!( event.data.provider_request_body_ref.as_deref(), Some("blob://provider-request-1") ); assert!(event .data .request_metadata .as_ref() .is_none_or(|value| { value.get("provider_request_body_ref").is_none() })); } #[test] fn stream_terminal_usage_records_base64_response_sizes_in_metadata() { let provider_bytes = b"{\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}"; let client_bytes = b"data: {\"id\":\"chatcmpl_123\"}\n\ndata: [DONE]\n"; let plan = ExecutionPlan { request_id: "req-stream-bytes-1".to_string(), candidate_id: Some("cand-stream-bytes-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:chat".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewayStreamReportRequest { trace_id: "trace-stream-bytes-1".to_string(), report_kind: "openai_chat_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:chat", "provider_api_format": "openai:responses", "trace_id": "trace-stream-bytes-1" })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: Some( base64::engine::general_purpose::STANDARD.encode(provider_bytes), ), provider_body_state: Some(UsageBodyCaptureState::Inline), client_body_base64: Some( base64::engine::general_purpose::STANDARD.encode(client_bytes), ), client_body_state: Some(UsageBodyCaptureState::Inline), terminal_summary: None, telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!( event .data .request_metadata .as_ref() .and_then(|value| value.get("provider_response_body_base64_bytes")) .and_then(|value| value.as_u64()), Some(provider_bytes.len() as u64) ); assert_eq!( event .data .request_metadata .as_ref() .and_then(|value| value.get("client_response_body_base64_bytes")) .and_then(|value| value.as_u64()), Some(client_bytes.len() as u64) ); } #[test] fn stream_terminal_usage_preserves_large_sse_response_for_repo_body_storage() { let large_delta = "x".repeat(MAX_USAGE_CAPTURE_BYTES); let sse_body = format!( concat!( "event: response.created\n", "data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}}}\n\n", "event: response.output_text.delta\n", "data: {{\"type\":\"response.output_text.delta\",\"delta\":\"{delta}\"}}\n\n", "event: response.completed\n", "data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{{\"input_tokens\":3,\"input_tokens_details\":{{\"cached_tokens\":2,\"cached_creation_tokens\":1}},\"output_tokens\":5,\"output_tokens_details\":{{\"reasoning_tokens\":1}},\"total_tokens\":8}}}}}}\n\n", "data: [DONE]\n", ), delta = large_delta ); let plan = ExecutionPlan { request_id: "req-stream-usage-large-1".to_string(), candidate_id: Some("cand-stream-usage-large-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody { json_body: None, body_bytes_b64: None, body_ref: None, }, stream: true, client_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewayStreamReportRequest { trace_id: "trace-stream-usage-large-1".to_string(), report_kind: "openai_responses_stream_success".to_string(), report_context: Some(json!({ "client_api_format": "openai:responses", "provider_api_format": "openai:responses", })), status_code: 200, headers: BTreeMap::new(), provider_body_base64: Some(base64::engine::general_purpose::STANDARD.encode(&sse_body)), provider_body_state: Some(UsageBodyCaptureState::Truncated), client_body_base64: None, client_body_state: Some(UsageBodyCaptureState::None), terminal_summary: None, telemetry: None, }; let event = build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.data.input_tokens, Some(3)); assert_eq!(event.data.output_tokens, Some(5)); assert_eq!(event.data.total_tokens, Some(11)); assert_eq!(event.data.cache_creation_input_tokens, Some(1)); assert_eq!(event.data.cache_read_input_tokens, Some(2)); assert_eq!( event .data .response_body .as_ref() .and_then(|value| value.get("chunks")) .and_then(Value::as_array) .and_then(|chunks| chunks.get(1)) .and_then(|value| value.get("delta")) .and_then(Value::as_str), Some(large_delta.as_str()) ); } #[test] fn sync_terminal_usage_does_not_fallback_request_body_to_plan_when_client_echo_is_absent() { let plan = ExecutionPlan { request_id: "req-sync-no-client-echo-1".to_string(), candidate_id: Some("cand-sync-no-client-echo-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({ "model": "gpt-5.4", "input": [{"role": "user", "content": "provider-side compiled body"}], })), stream: false, client_api_format: "claude:messages".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-no-client-echo-1".to_string(), report_kind: "claude_cli_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "claude:messages", "provider_api_format": "openai:responses", "needs_conversion": true, "original_request_body": null, })), status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({"id": "resp_1"})), client_body_json: None, body_base64: None, telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert!(event.data.request_body.is_none()); assert_eq!( event.data.provider_request_body, Some(json!({ "model": "gpt-5.4", "input": [{"role": "user", "content": "provider-side compiled body"}], })) ); } #[test] fn sync_terminal_usage_treats_null_error_field_as_success() { let plan = ExecutionPlan { request_id: "req-sync-null-error-1".to_string(), candidate_id: Some("cand-sync-null-error-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/messages".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"model": "gpt-5.4"})), stream: false, client_api_format: "claude:messages".to_string(), provider_api_format: "openai:responses".to_string(), model_name: Some("gpt-5.4".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let payload = GatewaySyncReportRequest { trace_id: "trace-sync-null-error-1".to_string(), report_kind: "claude_cli_sync_success".to_string(), report_context: Some(json!({ "client_api_format": "claude:messages", "provider_api_format": "openai:responses", "needs_conversion": true, })), status_code: 200, headers: BTreeMap::new(), body_json: Some(json!({ "id": "resp_1", "status": "completed", "error": null, "usage": { "input_tokens": 24, "output_tokens": 11, "total_tokens": 35 } })), client_body_json: Some(json!({ "type": "message", "usage": { "input_tokens": 24, "output_tokens": 11 } })), body_base64: None, telemetry: None, }; let event = build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) .expect("usage event should build"); assert_eq!(event.event_type, UsageEventType::Completed); assert_eq!(event.data.status_code, Some(200)); assert_eq!(event.data.input_tokens, Some(24)); assert_eq!(event.data.output_tokens, Some(11)); } #[test] fn manual_terminal_seed_event_builder_sanitizes_headers_and_metadata_but_preserves_bodies() { let event = build_terminal_usage_event_from_seed(TerminalUsageSeed { terminal_state: UsageTerminalState::Completed, client_contract: "openai:chat".to_string(), provider_contract: "openai:chat".to_string(), request_id: "req-manual-seed-1".to_string(), user_id: Some("user-1".to_string()), api_key_id: Some("key-1".to_string()), username: Some("alice".to_string()), api_key_name: Some("primary".to_string()), provider_name: "OpenAI".to_string(), model: "gpt-5".to_string(), target_model: None, model_id: None, global_model_id: None, provider_id: Some("provider-1".to_string()), provider_endpoint_id: Some("endpoint-1".to_string()), provider_api_key_id: Some("upstream-key-1".to_string()), request_type: "chat".to_string(), has_format_conversion: false, is_stream: false, body_refs: UsageBodyRefsSeed::default(), body_states: UsageBodyStatesSeed::default(), routing: UsageRoutingSeed { candidate_id: Some("cand-1".to_string()), ..UsageRoutingSeed::default() }, status_code: 200, response_time_ms: Some(123), first_byte_time_ms: Some(45), request_headers: Some(json!({ "authorization": "Bearer very-secret-token", "accept": "application/json" })), request_body: Some(json!({ "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) })), provider_request_headers: Some(json!({ "x-api-key": "sk-proj-super-secret" })), provider_request: None, provider_response_headers: Some(json!({ "set-cookie": "session=extremely-secret-cookie" })), provider_response: Some(json!({ "usage": { "input_tokens": 1, "output_tokens": 2, "total_tokens": 3 } })), client_response_headers: Some(json!({ "authorization": "Bearer client-secret" })), client_response: None, request_metadata: Some(json!({ "candidate_id": "cand-1", "billing_snapshot": { "payload": "x".repeat(32 * 1024) } })), audit_payload: None, standardized_usage: None, }) .expect("usage event should build"); assert_eq!( event.data.request_headers, Some(json!({ "authorization": "Bear****oken", "accept": "application/json" })) ); assert_eq!( event.data.provider_request_headers, Some(json!({ "x-api-key": "sk-p****cret" })) ); assert_eq!( event.data.response_headers, Some(json!({ "set-cookie": "sess****okie" })) ); assert_eq!( event.data.client_response_headers, Some(json!({ "authorization": "Bear****cret" })) ); assert_eq!( event.data.request_body, Some(json!({ "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) })) ); assert_eq!(event.data.input_tokens, Some(1)); assert_eq!(event.data.output_tokens, Some(2)); assert_eq!(event.data.total_tokens, Some(3)); assert_eq!(event.data.candidate_id.as_deref(), Some("cand-1")); assert_eq!( event .data .request_metadata .as_ref() .and_then(|value| value.get("candidate_id")), None ); assert_eq!( event .data .request_metadata .as_ref() .and_then(|value| value.get("billing_snapshot")), Some(&json!({ "truncated": true, "reason": "usage_request_metadata_limits_exceeded", "max_depth": 32, "max_nodes": 4_000, "max_bytes": 16 * 1024, "value_kind": "object" })) ); } #[test] fn lifecycle_seed_builder_sanitizes_manual_request_metadata() { let record = build_pending_usage_record_from_seed( &LifecycleUsageSeed { request_id: "req-lifecycle-manual-1".to_string(), user_id: Some("user-1".to_string()), api_key_id: Some("key-1".to_string()), username: Some("alice".to_string()), api_key_name: Some("primary".to_string()), provider_name: "OpenAI".to_string(), model: "gpt-5".to_string(), target_model: None, model_id: None, global_model_id: None, provider_id: Some("provider-1".to_string()), provider_endpoint_id: Some("endpoint-1".to_string()), provider_api_key_id: Some("upstream-key-1".to_string()), request_type: "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: false, body_states: UsageBodyStatesSeed::default(), routing: UsageRoutingSeed { candidate_id: Some("cand-1".to_string()), ..UsageRoutingSeed::default() }, request_metadata: Some(json!({ "billing_snapshot": { "payload": "x".repeat(32 * 1024) } })), }, 1_700_000_000, ) .expect("pending record should build"); assert_eq!(record.candidate_id.as_deref(), Some("cand-1")); assert_eq!( record.request_metadata, Some(json!({ "billing_snapshot": { "truncated": true, "reason": "usage_request_metadata_limits_exceeded", "max_depth": 32, "max_nodes": 4_000, "max_bytes": 16 * 1024, "value_kind": "object" } })) ); } #[test] fn usage_event_data_seed_masks_headers_before_outcome_paths_use_it() { let plan = ExecutionPlan { request_id: "req-seed-sanitize-1".to_string(), candidate_id: Some("cand-seed-sanitize-1".to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), url: "https://example.com/v1/chat/completions".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: RequestBody::from_json(json!({"model": "gpt-5"})), stream: false, client_api_format: "openai:chat".to_string(), provider_api_format: "openai:chat".to_string(), model_name: Some("gpt-5".to_string()), proxy: None, transport_profile: None, timeouts: None, }; let data = build_usage_event_data_seed( &plan, Some(&json!({ "original_headers": { "authorization": "Bearer outcome-secret", "accept": "application/json" }, "original_request_body": { "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) } })), ); assert_eq!( data.request_headers, Some(json!({ "authorization": "Bear****cret", "accept": "application/json" })) ); assert_eq!( data.request_body, Some(json!({ "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) })) ); } #[test] fn masks_known_sensitive_header_values() { let token = "Bearer eyJhbGciOiJSUzI1NiJ9.payload-here.signature-tail"; let masked = mask_header_value("authorization", token); assert!(masked.starts_with("Bear")); assert!(masked.ends_with("tail")); assert!(masked.contains("****")); assert!(!masked.contains("payload-here")); // 大小写不敏感 assert_eq!( mask_header_value("Authorization", "12345678"), "****", "短值整体替换为 ****", ); assert_eq!(mask_header_value("X-Api-Key", "abcdefghij"), "abcd****ghij",); // 非敏感头保持原样 assert_eq!( mask_header_value("user-agent", "codex-tui/0.1"), "codex-tui/0.1", ); } #[test] fn headers_to_json_masks_sensitive_headers_at_source() { let mut headers = BTreeMap::new(); headers.insert( "authorization".to_string(), "Bearer eyJhbGciOiJSUzI1NiJ9.body.signature".to_string(), ); headers.insert("user-agent".to_string(), "codex-tui/0.1".to_string()); headers.insert( "x-api-key".to_string(), "sk-proj-1234567890abcdef".to_string(), ); let value = headers_to_json(&headers).expect("expected object"); let object = value.as_object().expect("expected object"); let auth = object .get("authorization") .and_then(|v| v.as_str()) .expect("authorization should be string"); assert!(auth.starts_with("Bear")); assert!(auth.contains("****")); assert!(!auth.contains("eyJhbGciOiJSUzI1NiJ9")); let api_key = object .get("x-api-key") .and_then(|v| v.as_str()) .expect("x-api-key should be string"); assert!(api_key.starts_with("sk-p")); assert!(api_key.contains("****")); assert!(!api_key.contains("1234567890")); assert_eq!( object.get("user-agent").and_then(|v| v.as_str()), Some("codex-tui/0.1"), ); } #[test] fn headers_to_json_returns_none_for_empty_headers() { assert!(headers_to_json(&BTreeMap::new()).is_none()); } #[test] fn mask_sensitive_headers_in_json_value_handles_object_form() { let value = json!({ "Authorization": "Bearer eyJhbGciOiJSUzI1NiJ9.body.signature", "Cookie": "session=verylongcookievalue1234", "Accept": "application/json", }); let masked = mask_sensitive_headers_in_json_value(Some(value)).expect("masked value should exist"); let object = masked.as_object().expect("expected object"); let auth = object .get("Authorization") .and_then(|v| v.as_str()) .expect("Authorization should be string"); assert!(auth.contains("****")); assert!(!auth.contains("eyJhbGciOiJSUzI1NiJ9")); let cookie = object .get("Cookie") .and_then(|v| v.as_str()) .expect("Cookie should be string"); assert!(cookie.contains("****")); assert!(!cookie.contains("verylongcookievalue")); assert_eq!( object.get("Accept").and_then(|v| v.as_str()), Some("application/json"), ); } #[test] fn mask_sensitive_headers_passthrough_for_non_object() { // None 输入返回 None assert!(mask_sensitive_headers_in_json_value(None).is_none()); // 非 object 输入原样返回 let masked = mask_sensitive_headers_in_json_value(Some(json!("not an object"))); assert_eq!(masked, Some(json!("not an object"))); } #[test] fn resolve_error_message_extracts_generic_message_from_base64_body() { let body_base64 = base64::engine::general_purpose::STANDARD.encode( serde_json::to_vec(&json!({ "message": "upstream exploded" })) .expect("json should serialize"), ); assert_eq!( resolve_error_message(500, None, Some(body_base64.as_str())), Some("upstream exploded".to_string()), ); } #[test] fn parse_sse_body_for_storage_handles_crlf_and_cr_line_endings() { let sse_body = concat!( "data: {\"id\":\"chunk-1\"}\r\n", "\r\n", "data: {\"id\":\"chunk-2\"}\r", "\r", "data: [DONE]\r", ); assert_eq!( parse_sse_body_for_storage(sse_body), Some(json!({ "chunks": [ {"id": "chunk-1"}, {"id": "chunk-2"} ], "metadata": { "stream": true, "total_chunks": 2, "stored_chunks": 2, "content_length": sse_body.len(), "has_completion": true } })), ); } #[test] fn extract_token_counts_from_value_handles_crlf_and_cr_sse_text() { let sse_body = concat!( "event: response.created\r\n", "data: {\"type\":\"response.created\"}\r\n", "\r\n", "event: response.completed\r", "data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":2,\"cached_creation_tokens\":1},\"output_tokens\":5,\"total_tokens\":8}}}\r", "\r", "data: [DONE]\r", ); assert_eq!( extract_token_counts_from_value(&Value::String(sse_body.to_string())), Some((3, 5, 11)), ); } }