Aether/crates/aether-usage-runtime/src/write.rs

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<String>,
    candidate_index: Option<u64>,
    key_name: Option<String>,
    planner_kind: Option<String>,
    route_family: Option<String>,
    route_kind: Option<String>,
    execution_path: Option<String>,
    local_execution_runtime_miss_reason: Option<String>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageBodyRefsSeed {
    request_body_ref: Option<String>,
    provider_request_body_ref: Option<String>,
    response_body_ref: Option<String>,
    client_response_body_ref: Option<String>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageBodyStatesSeed {
    request_body_state: Option<UsageBodyCaptureState>,
    provider_request_body_state: Option<UsageBodyCaptureState>,
    response_body_state: Option<UsageBodyCaptureState>,
    client_response_body_state: Option<UsageBodyCaptureState>,
}

#[derive(Debug, Clone)]
struct RuntimeRequestCaptureSeed {
    request_body: Option<Value>,
    request_body_ref: Option<String>,
    provider_request: Option<Value>,
    provider_request_body_ref: Option<String>,
    body_states: UsageBodyStatesSeed,
}

#[derive(Debug, Clone, PartialEq)]
pub struct LifecycleUsageSeed {
    pub request_id: String,
    pub user_id: Option<String>,
    pub api_key_id: Option<String>,
    pub username: Option<String>,
    pub api_key_name: Option<String>,
    pub provider_name: String,
    pub model: String,
    pub target_model: Option<String>,
    pub model_id: Option<String>,
    pub global_model_id: Option<String>,
    pub provider_id: Option<String>,
    pub provider_endpoint_id: Option<String>,
    pub provider_api_key_id: Option<String>,
    pub request_type: String,
    pub api_format: Option<String>,
    pub api_family: Option<String>,
    pub endpoint_kind: Option<String>,
    pub endpoint_api_format: Option<String>,
    pub provider_api_family: Option<String>,
    pub provider_endpoint_kind: Option<String>,
    pub has_format_conversion: Option<bool>,
    pub is_stream: bool,
    routing: UsageRoutingSeed,
    body_states: UsageBodyStatesSeed,
    pub request_metadata: Option<Value>,
}

#[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<String>,
    pub api_key_id: Option<String>,
    pub username: Option<String>,
    pub api_key_name: Option<String>,
    pub provider_name: String,
    pub model: String,
    pub target_model: Option<String>,
    pub model_id: Option<String>,
    pub global_model_id: Option<String>,
    pub provider_id: Option<String>,
    pub provider_endpoint_id: Option<String>,
    pub provider_api_key_id: Option<String>,
    pub request_type: String,
    pub has_format_conversion: bool,
    pub is_stream: bool,
    pub request_headers: Option<Value>,
    pub request_body: Option<Value>,
    pub provider_request_headers: Option<Value>,
    pub provider_request: Option<Value>,
    body_refs: UsageBodyRefsSeed,
    body_states: UsageBodyStatesSeed,
    routing: UsageRoutingSeed,
    pub request_metadata: Option<Value>,
}

#[derive(Debug, Clone)]
pub struct SyncTerminalUsagePayloadSeed {
    pub report_kind: String,
    pub status_code: u16,
    pub response_time_ms: Option<u64>,
    pub first_byte_time_ms: Option<u64>,
    pub provider_response_headers: Option<Value>,
    pub client_response_headers: Option<Value>,
    pub provider_response_full: Option<Value>,
    pub provider_response_body_state: Option<UsageBodyCaptureState>,
    pub client_response: Option<Value>,
    pub client_response_body_state: Option<UsageBodyCaptureState>,
    pub capture_metadata: Option<Value>,
}

#[derive(Debug, Clone)]
pub struct StreamTerminalUsagePayloadSeed {
    pub report_kind: String,
    pub status_code: u16,
    pub response_time_ms: Option<u64>,
    pub first_byte_time_ms: Option<u64>,
    pub provider_response_headers: Option<Value>,
    pub client_response_headers: Option<Value>,
    pub provider_response_full: Option<Value>,
    pub provider_response_body_state: Option<UsageBodyCaptureState>,
    pub client_response: Option<Value>,
    pub client_response_body_state: Option<UsageBodyCaptureState>,
    pub standardized_usage: Option<StandardizedUsage>,
    pub capture_metadata: Option<Value>,
}

#[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<String>,
    pub api_key_id: Option<String>,
    pub username: Option<String>,
    pub api_key_name: Option<String>,
    pub provider_name: String,
    pub model: String,
    pub target_model: Option<String>,
    pub model_id: Option<String>,
    pub global_model_id: Option<String>,
    pub provider_id: Option<String>,
    pub provider_endpoint_id: Option<String>,
    pub provider_api_key_id: Option<String>,
    pub request_type: String,
    pub has_format_conversion: bool,
    pub is_stream: bool,
    pub status_code: u16,
    pub response_time_ms: Option<u64>,
    pub first_byte_time_ms: Option<u64>,
    pub request_headers: Option<Value>,
    pub request_body: Option<Value>,
    pub provider_request_headers: Option<Value>,
    pub provider_request: Option<Value>,
    body_refs: UsageBodyRefsSeed,
    body_states: UsageBodyStatesSeed,
    pub provider_response_headers: Option<Value>,
    pub provider_response: Option<Value>,
    pub client_response_headers: Option<Value>,
    pub client_response: Option<Value>,
    routing: UsageRoutingSeed,
    pub request_metadata: Option<Value>,
    pub audit_payload: Option<Value>,
    pub standardized_usage: Option<StandardizedUsage>,
}

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<u16>,
    response_time_ms: Option<u64>,
    first_byte_time_ms: Option<u64>,
    response_headers: Option<Value>,
    client_response_headers: Option<Value>,
    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<UpsertUsageRecord, DataLayerError> {
    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<UpsertUsageRecord, DataLayerError> {
    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<UpsertUsageRecord, DataLayerError> {
    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<UpsertUsageRecord, DataLayerError> {
    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<UpsertUsageRecord, DataLayerError> {
    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<ExecutionTelemetry>,
    updated_at_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
    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<UsageEvent, DataLayerError> {
    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<UsageEvent, DataLayerError> {
    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<UsageEvent, DataLayerError> {
    build_terminal_usage_event_from_seed(outcome)
}

pub fn build_terminal_usage_event_from_seed(
    seed: TerminalUsageSeed,
) -> Result<UsageEvent, DataLayerError> {
    build_terminal_usage_event_from_seed_impl(seed, false)
}

fn build_terminal_usage_event_from_seed_impl(
    seed: TerminalUsageSeed,
    trusted_request_metadata: bool,
) -> Result<UsageEvent, DataLayerError> {
    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<String, Value>>,
    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<UpsertUsageRecord, DataLayerError> {
    let LifecycleUsageRecordInput { seed, options } = input;
    build_lifecycle_usage_record_impl(seed, options)
}

fn build_lifecycle_usage_record_owned(
    input: OwnedLifecycleUsageRecordInput,
) -> Result<UpsertUsageRecord, DataLayerError> {
    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<UpsertUsageRecord, DataLayerError> {
    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<String, Value>>,
) -> 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<String> {
    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<String, Value>>, key: &str) -> Option<String> {
    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<String, Value>>, key: &str) -> Option<u64> {
    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<u64> {
    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<String, Value>>, key: &str) -> Option<bool> {
    context
        .and_then(|value| value.get(key))
        .and_then(Value::as_bool)
}

fn context_value_ref<'a>(context: Option<&'a Map<String, Value>>, key: &str) -> Option<&'a Value> {
    context.and_then(|value| value.get(key))
}

fn context_value(context: Option<&Map<String, Value>>, key: &str) -> Option<Value> {
    context.and_then(|value| value.get(key)).cloned()
}

fn context_usage_value(context: Option<&Map<String, Value>>, key: &str) -> Option<Value> {
    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<String, Value>>, key: &str) -> Option<Value> {
    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<String, Value>>, 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<String, Value>>,
) -> 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<String, Value>>,
) -> Option<Value> {
    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<String, Value>>,
    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<Value> {
    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<String, Value>>,
) -> 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<Value> {
    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> {
    value.cloned().map(capture_usage_storage_value)
}

fn clone_usage_body_value(value: Option<&Value>) -> Option<Value> {
    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<Value>) -> Option<Value> {
    value.map(capture_usage_storage_value)
}

fn sanitize_usage_header_capture(value: Option<Value>) -> Option<Value> {
    mask_sensitive_headers_in_json_value(value).map(capture_usage_storage_value)
}

fn trim_owned_non_empty_string(value: String) -> Option<String> {
    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<String>) -> Option<String> {
    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::<usize>(),
                    );
                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<String> {
    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<String, String>) -> Option<Value> {
    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::<Vec<_>>()
        .into_iter()
        .rev()
        .collect();
    format!("{prefix}****{suffix}")
}

/// 对 JSON 形式的 headers 做就地脱敏。仅当 value 是 Object 时才会处理；
/// 其它形式的值保持不变。
fn mask_sensitive_headers_in_json_value(value: Option<Value>) -> Option<Value> {
    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<String> {
    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<String> {
    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<String> {
    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<String> {
    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<Value> {
    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::<Value>(&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<F>(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<F>(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<Value> {
    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::<Value>(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::<Value>(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<String, Value>,
    usage_details: Option<&serde_json::Map<String, Value>>,
) -> 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<String, Value>,
    usage_details: Option<&serde_json::Map<String, Value>>,
) -> 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<String>) -> Option<String> {
    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)),
        );
    }
}