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

use std::collections::BTreeMap;

use aether_billing::map_usage_from_response;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data::repository::usage::UpsertUsageRecord;
use aether_data::DataLayerError;
use base64::Engine as _;
use serde_json::{Map, Value};

use super::reporting::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use super::{UsageEvent, UsageEventData, UsageEventType};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum UsageLifecycleState {
    Pending,
    Streaming,
    Completed,
    Failed,
    Cancelled,
}

pub(crate) fn build_pending_usage_record(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    updated_at_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
    build_upsert_usage_record(
        plan,
        report_context,
        UsageEventData::default(),
        UsageLifecycleState::Pending,
        None,
        updated_at_unix_secs,
    )
}

pub(crate) fn build_streaming_usage_record(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    status_code: u16,
    response_headers: &BTreeMap<String, String>,
    telemetry: Option<&ExecutionTelemetry>,
    updated_at_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
    build_upsert_usage_record(
        plan,
        report_context,
        UsageEventData {
            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: Some(headers_to_json(response_headers)),
            client_response_headers: Some(headers_to_json(response_headers)),
            ..UsageEventData::default()
        },
        UsageLifecycleState::Streaming,
        None,
        updated_at_unix_secs,
    )
}

pub(crate) fn build_sync_terminal_usage_event(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    payload: &GatewaySyncReportRequest,
    event_type: UsageEventType,
) -> Result<UsageEvent, DataLayerError> {
    let mut data = UsageEventData {
        status_code: Some(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),
        response_headers: Some(headers_to_json(&payload.headers)),
        client_response_headers: Some(headers_to_json(&payload.headers)),
        response_body: payload
            .body_json
            .clone()
            .or_else(|| decode_body_for_storage(payload.body_base64.as_deref())),
        client_response_body: payload.client_body_json.clone(),
        error_message: resolve_error_message(
            payload.status_code,
            payload.body_json.as_ref(),
            payload.body_base64.as_deref(),
        ),
        error_category: resolve_error_category(payload.status_code, event_type),
        ..UsageEventData::default()
    };

    let standardized_response_body = payload
        .body_json
        .as_ref()
        .cloned()
        .or_else(|| data.response_body.clone());
    if let Some(tokens) = standardized_response_body.as_ref() {
        apply_standardized_usage(
            resolve_billing_api_format(plan, report_context),
            tokens,
            &mut data,
        );
    }
    if data.total_tokens.is_none() {
        if let Some(tokens) = payload
            .body_json
            .as_ref()
            .and_then(extract_token_counts_from_json)
            .or_else(|| {
                payload
                    .body_base64
                    .as_deref()
                    .and_then(extract_token_counts_from_base64_payload)
            })
        {
            data.input_tokens = Some(tokens.0);
            data.output_tokens = Some(tokens.1);
            data.total_tokens = Some(tokens.2);
        }
    }

    build_usage_event(plan, report_context, data, event_type)
}

pub(crate) fn build_stream_terminal_usage_event(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    payload: &GatewayStreamReportRequest,
    event_type: UsageEventType,
) -> Result<UsageEvent, DataLayerError> {
    let mut data = UsageEventData {
        status_code: Some(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),
        response_headers: Some(headers_to_json(&payload.headers)),
        client_response_headers: Some(headers_to_json(&payload.headers)),
        response_body: decode_body_for_storage(payload.body_base64.as_deref()),
        error_message: resolve_error_message(
            payload.status_code,
            None,
            payload.body_base64.as_deref(),
        ),
        error_category: resolve_error_category(payload.status_code, event_type),
        ..UsageEventData::default()
    };

    let standardized_response_body = data.response_body.clone();
    if let Some(response_body) = standardized_response_body.as_ref() {
        apply_standardized_usage(
            resolve_billing_api_format(plan, report_context),
            response_body,
            &mut data,
        );
    }
    if data.total_tokens.is_none() {
        if let Some(tokens) = payload
            .body_base64
            .as_deref()
            .and_then(extract_token_counts_from_base64_payload)
        {
            data.input_tokens = Some(tokens.0);
            data.output_tokens = Some(tokens.1);
            data.total_tokens = Some(tokens.2);
        }
    }

    build_usage_event(plan, report_context, data, event_type)
}

pub(crate) fn build_upsert_usage_record_from_event(
    event: &UsageEvent,
) -> Result<UpsertUsageRecord, DataLayerError> {
    let lifecycle_state = match event.event_type {
        UsageEventType::Pending => UsageLifecycleState::Pending,
        UsageEventType::Streaming => UsageLifecycleState::Streaming,
        UsageEventType::Completed => UsageLifecycleState::Completed,
        UsageEventType::Failed => UsageLifecycleState::Failed,
        UsageEventType::Cancelled => UsageLifecycleState::Cancelled,
    };

    let now_unix_secs = event.timestamp_ms / 1_000;
    build_upsert_usage_record(
        &ExecutionPlan {
            request_id: event.request_id.clone(),
            candidate_id: None,
            provider_name: Some(event.data.provider_name.clone()),
            provider_id: event.data.provider_id.clone().unwrap_or_default(),
            endpoint_id: event.data.provider_endpoint_id.clone().unwrap_or_default(),
            key_id: event.data.provider_api_key_id.clone().unwrap_or_default(),
            method: "POST".to_string(),
            url: String::new(),
            headers: BTreeMap::new(),
            content_type: None,
            content_encoding: None,
            body: aether_contracts::RequestBody {
                json_body: None,
                body_bytes_b64: None,
                body_ref: None,
            },
            stream: event.data.is_stream.unwrap_or(false),
            client_api_format: event.data.api_format.clone().unwrap_or_default(),
            provider_api_format: event
                .data
                .endpoint_api_format
                .clone()
                .or_else(|| event.data.api_format.clone())
                .unwrap_or_default(),
            model_name: Some(event.data.model.clone()),
            proxy: None,
            tls_profile: None,
            timeouts: None,
        },
        None,
        event.data.clone(),
        lifecycle_state,
        Some(now_unix_secs),
        now_unix_secs,
    )
}

fn build_usage_event(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    data: UsageEventData,
    event_type: UsageEventType,
) -> Result<UsageEvent, DataLayerError> {
    let base = build_base_usage_data(plan, report_context);
    Ok(UsageEvent::new(
        event_type,
        plan.request_id.clone(),
        merge_usage_data(base, data),
    ))
}

fn build_upsert_usage_record(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
    data: UsageEventData,
    lifecycle_state: UsageLifecycleState,
    finalized_at_unix_secs: Option<u64>,
    updated_at_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
    let data = merge_usage_data(build_base_usage_data(plan, report_context), data);
    let (status, billing_status) = lifecycle_status_and_billing(lifecycle_state);

    Ok(UpsertUsageRecord {
        request_id: plan.request_id.clone(),
        user_id: data.user_id,
        api_key_id: data.api_key_id,
        username: data.username,
        api_key_name: data.api_key_name,
        provider_name: data.provider_name,
        model: data.model,
        target_model: data.target_model,
        provider_id: empty_to_none(data.provider_id),
        provider_endpoint_id: empty_to_none(data.provider_endpoint_id),
        provider_api_key_id: empty_to_none(data.provider_api_key_id),
        request_type: data.request_type,
        api_format: data.api_format,
        api_family: data.api_family,
        endpoint_kind: data.endpoint_kind,
        endpoint_api_format: data.endpoint_api_format,
        provider_api_family: data.provider_api_family,
        provider_endpoint_kind: data.provider_endpoint_kind,
        has_format_conversion: data.has_format_conversion,
        is_stream: data.is_stream,
        input_tokens: data.input_tokens,
        output_tokens: data.output_tokens,
        total_tokens: data.total_tokens,
        cache_creation_input_tokens: data.cache_creation_input_tokens,
        cache_read_input_tokens: data.cache_read_input_tokens,
        cache_creation_cost_usd: data.cache_creation_cost_usd,
        cache_read_cost_usd: data.cache_read_cost_usd,
        output_price_per_1m: data.output_price_per_1m,
        total_cost_usd: data.total_cost_usd,
        actual_total_cost_usd: data.actual_total_cost_usd,
        status_code: data.status_code,
        error_message: data.error_message,
        error_category: data.error_category,
        response_time_ms: data.response_time_ms,
        first_byte_time_ms: data.first_byte_time_ms,
        status: status.to_string(),
        billing_status: billing_status.to_string(),
        request_headers: data.request_headers,
        request_body: data.request_body,
        provider_request_headers: data.provider_request_headers,
        provider_request_body: data.provider_request_body,
        response_headers: data.response_headers,
        response_body: data.response_body,
        client_response_headers: data.client_response_headers,
        client_response_body: data.client_response_body,
        request_metadata: data.request_metadata,
        finalized_at_unix_secs,
        created_at_unix_secs: Some(updated_at_unix_secs),
        updated_at_unix_secs,
    })
}

fn build_base_usage_data(plan: &ExecutionPlan, report_context: Option<&Value>) -> UsageEventData {
    let context = report_context.and_then(Value::as_object);
    let api_format = context_string(context, "client_api_format")
        .or_else(|| non_empty_string(Some(plan.client_api_format.clone())));
    let endpoint_api_format = context_string(context, "provider_api_format")
        .or_else(|| non_empty_string(Some(plan.provider_api_format.clone())));
    let model = context_string(context, "model")
        .or_else(|| plan.model_name.clone())
        .unwrap_or_else(|| "unknown".to_string());
    let provider_name = context_string(context, "provider_name")
        .or_else(|| plan.provider_name.clone())
        .filter(|value| !value.trim().is_empty())
        .unwrap_or_else(|| "unknown".to_string());

    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"),
        provider_id: context_string(context, "provider_id")
            .or_else(|| non_empty_string(Some(plan.provider_id.clone()))),
        provider_endpoint_id: context_string(context, "endpoint_id")
            .or_else(|| non_empty_string(Some(plan.endpoint_id.clone()))),
        provider_api_key_id: context_string(context, "key_id")
            .or_else(|| non_empty_string(Some(plan.key_id.clone()))),
        request_type: Some(infer_request_type(api_format.as_deref())),
        api_format: api_format.clone(),
        api_family: api_format
            .as_deref()
            .and_then(infer_api_family)
            .map(ToOwned::to_owned),
        endpoint_kind: api_format
            .as_deref()
            .and_then(infer_endpoint_kind)
            .map(ToOwned::to_owned),
        endpoint_api_format: endpoint_api_format.clone(),
        provider_api_family: endpoint_api_format
            .as_deref()
            .and_then(infer_api_family)
            .map(ToOwned::to_owned),
        provider_endpoint_kind: endpoint_api_format
            .as_deref()
            .and_then(infer_endpoint_kind)
            .map(ToOwned::to_owned),
        has_format_conversion: context_bool(context, "needs_conversion"),
        is_stream: Some(plan.stream),
        request_headers: context_value(context, "original_headers"),
        request_body: context_value(context, "original_request_body")
            .or_else(|| plan.body.json_body.clone()),
        provider_request_headers: context_value(context, "provider_request_headers")
            .or_else(|| Some(headers_to_json(&plan.headers))),
        provider_request_body: context_value(context, "provider_request_body")
            .or_else(|| plan.body.json_body.clone()),
        request_metadata: Some(Value::Object(Map::from_iter([
            (
                "request_id".to_string(),
                Value::String(plan.request_id.clone()),
            ),
            (
                "candidate_id".to_string(),
                plan.candidate_id
                    .clone()
                    .map(Value::String)
                    .unwrap_or(Value::Null),
            ),
        ]))),
        ..UsageEventData::default()
    }
}

fn merge_usage_data(base: UsageEventData, override_data: UsageEventData) -> UsageEventData {
    UsageEventData {
        user_id: override_data.user_id.or(base.user_id),
        api_key_id: override_data.api_key_id.or(base.api_key_id),
        username: override_data.username.or(base.username),
        api_key_name: override_data.api_key_name.or(base.api_key_name),
        provider_name: if override_data.provider_name.trim().is_empty() {
            base.provider_name
        } else {
            override_data.provider_name
        },
        model: if override_data.model.trim().is_empty() {
            base.model
        } else {
            override_data.model
        },
        target_model: override_data.target_model.or(base.target_model),
        provider_id: override_data.provider_id.or(base.provider_id),
        provider_endpoint_id: override_data
            .provider_endpoint_id
            .or(base.provider_endpoint_id),
        provider_api_key_id: override_data
            .provider_api_key_id
            .or(base.provider_api_key_id),
        request_type: override_data.request_type.or(base.request_type),
        api_format: override_data.api_format.or(base.api_format),
        api_family: override_data.api_family.or(base.api_family),
        endpoint_kind: override_data.endpoint_kind.or(base.endpoint_kind),
        endpoint_api_format: override_data
            .endpoint_api_format
            .or(base.endpoint_api_format),
        provider_api_family: override_data
            .provider_api_family
            .or(base.provider_api_family),
        provider_endpoint_kind: override_data
            .provider_endpoint_kind
            .or(base.provider_endpoint_kind),
        has_format_conversion: override_data
            .has_format_conversion
            .or(base.has_format_conversion),
        is_stream: override_data.is_stream.or(base.is_stream),
        input_tokens: override_data.input_tokens.or(base.input_tokens),
        output_tokens: override_data.output_tokens.or(base.output_tokens),
        total_tokens: override_data.total_tokens.or(base.total_tokens),
        cache_creation_input_tokens: override_data
            .cache_creation_input_tokens
            .or(base.cache_creation_input_tokens),
        cache_read_input_tokens: override_data
            .cache_read_input_tokens
            .or(base.cache_read_input_tokens),
        cache_creation_cost_usd: override_data
            .cache_creation_cost_usd
            .or(base.cache_creation_cost_usd),
        cache_read_cost_usd: override_data
            .cache_read_cost_usd
            .or(base.cache_read_cost_usd),
        output_price_per_1m: override_data
            .output_price_per_1m
            .or(base.output_price_per_1m),
        total_cost_usd: override_data.total_cost_usd.or(base.total_cost_usd),
        actual_total_cost_usd: override_data
            .actual_total_cost_usd
            .or(base.actual_total_cost_usd),
        status_code: override_data.status_code.or(base.status_code),
        error_message: override_data.error_message.or(base.error_message),
        error_category: override_data.error_category.or(base.error_category),
        response_time_ms: override_data.response_time_ms.or(base.response_time_ms),
        first_byte_time_ms: override_data.first_byte_time_ms.or(base.first_byte_time_ms),
        request_headers: override_data.request_headers.or(base.request_headers),
        request_body: override_data.request_body.or(base.request_body),
        provider_request_headers: override_data
            .provider_request_headers
            .or(base.provider_request_headers),
        provider_request_body: override_data
            .provider_request_body
            .or(base.provider_request_body),
        response_headers: override_data.response_headers.or(base.response_headers),
        response_body: override_data.response_body.or(base.response_body),
        client_response_headers: override_data
            .client_response_headers
            .or(base.client_response_headers),
        client_response_body: override_data
            .client_response_body
            .or(base.client_response_body),
        request_metadata: merge_json_value(base.request_metadata, override_data.request_metadata),
    }
}

fn merge_json_value(base: Option<Value>, override_value: Option<Value>) -> Option<Value> {
    match (base, override_value) {
        (Some(Value::Object(mut base)), Some(Value::Object(override_object))) => {
            for (key, value) in override_object {
                base.insert(key, value);
            }
            Some(Value::Object(base))
        }
        (Some(base), None) => Some(base),
        (_, Some(override_value)) => Some(override_value),
        (None, None) => None,
    }
}

fn lifecycle_status_and_billing(state: UsageLifecycleState) -> (&'static str, &'static str) {
    match state {
        UsageLifecycleState::Pending => ("pending", "pending"),
        UsageLifecycleState::Streaming => ("streaming", "pending"),
        UsageLifecycleState::Completed => ("completed", "pending"),
        UsageLifecycleState::Failed => ("failed", "void"),
        UsageLifecycleState::Cancelled => ("cancelled", "void"),
    }
}

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_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(context: Option<&Map<String, Value>>, key: &str) -> Option<Value> {
    context.and_then(|value| value.get(key)).cloned()
}

fn non_empty_string(value: Option<String>) -> Option<String> {
    value
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty())
}

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 resolve_billing_api_format(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
) -> Option<String> {
    let context = report_context.and_then(Value::as_object);
    context_string(context, "provider_api_format")
        .or_else(|| non_empty_string(Some(plan.provider_api_format.clone())))
        .or_else(|| context_string(context, "client_api_format"))
        .or_else(|| non_empty_string(Some(plan.client_api_format.clone())))
}

fn apply_standardized_usage(
    api_format: Option<String>,
    response_body: &Value,
    data: &mut UsageEventData,
) {
    let Some(api_format) = api_format.as_deref() else {
        return;
    };
    if !response_body.is_object() {
        return;
    }
    let usage = map_usage_from_response(response_body, api_format);
    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_read_tokens > 0 {
        data.cache_read_input_tokens = Some(usage.cache_read_tokens as u64);
    }
    let total_tokens = usage
        .input_tokens
        .saturating_add(usage.output_tokens)
        .max(0) as u64;
    if total_tokens > 0 {
        data.total_tokens = Some(total_tokens);
    }
}

fn headers_to_json(headers: &BTreeMap<String, String>) -> Value {
    Value::Object(Map::from_iter(
        headers
            .iter()
            .map(|(key, value)| (key.clone(), Value::String(value.clone()))),
    ))
}

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()),
        _ => None,
    }
}

fn resolve_error_message(
    status_code: u16,
    body_json: Option<&Value>,
    body_base64: Option<&str>,
) -> Option<String> {
    let explicit_error_message = body_json
        .and_then(extract_explicit_error_message_from_json)
        .or_else(|| {
            body_base64
                .and_then(|value| decode_body_for_storage(Some(value)))
                .as_ref()
                .and_then(extract_explicit_error_message_from_json)
        });
    if explicit_error_message.is_some() {
        return explicit_error_message;
    }
    if status_code < 400 {
        return None;
    }

    body_json
        .and_then(extract_generic_error_message_from_json)
        .or_else(|| {
            body_base64
                .and_then(|value| decode_body_for_storage(Some(value)))
                .as_ref()
                .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) {
        return Some(Value::String(text));
    }
    Some(Value::String(body_base64.to_string()))
}

fn extract_token_counts_from_base64_payload(body_base64: &str) -> Option<(u64, u64, u64)> {
    let bytes = base64::engine::general_purpose::STANDARD
        .decode(body_base64)
        .ok()?;
    if let Ok(json_body) = serde_json::from_slice::<Value>(&bytes) {
        return extract_token_counts_from_json(&json_body);
    }
    let text = String::from_utf8(bytes).ok()?;
    extract_token_counts_from_sse_text(&text)
}

fn extract_token_counts_from_sse_text(text: &str) -> Option<(u64, u64, u64)> {
    let mut last_seen = None;
    for line in text.lines() {
        let line = line.trim();
        if line.is_empty() || !line.starts_with("data:") {
            continue;
        }
        let payload = line.trim_start_matches("data:").trim();
        if payload.is_empty() || payload == "[DONE]" {
            continue;
        }
        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_json(value: &Value) -> Option<(u64, u64, u64)> {
    if let Some(usage) = value.get("usage").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 total = usage
            .get("total_tokens")
            .and_then(Value::as_u64)
            .unwrap_or(input + output);
        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 total = usage
            .get("totalTokenCount")
            .and_then(Value::as_u64)
            .unwrap_or(input + output);
        return Some((input, output, total));
    }

    if let Some(response) = value.get("response") {
        return extract_token_counts_from_json(response);
    }

    None
}

fn empty_to_none(value: Option<String>) -> Option<String> {
    value
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty())
}

#[cfg(test)]
mod tests {
    use super::{build_upsert_usage_record_from_event, extract_token_counts_from_json};
    use crate::gateway::usage::{UsageEvent, UsageEventData, UsageEventType};
    use serde_json::json;

    #[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 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));
    }
}