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

use aether_ai_formats::api::{
    sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string,
};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::ExecutionPlan;
use aether_data_contracts::repository::usage::{
    extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body,
    PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
};
use serde_json::{json, Map, Value};

const MAX_USAGE_REQUEST_METADATA_DEPTH: usize = 32;
const MAX_USAGE_REQUEST_METADATA_NODES: usize = 4_000;
const MAX_USAGE_REQUEST_METADATA_BYTES: usize = 16 * 1024;
const MAX_USAGE_REQUEST_METADATA_STRING_BYTES: usize = 1_024;

pub(crate) fn build_usage_request_metadata_seed(
    _plan: &ExecutionPlan,
    context: Option<&Map<String, Value>>,
) -> Option<Value> {
    let mut metadata = Map::new();
    if let Some(context) = context {
        copy_allowed_metadata_fields(context, &mut metadata);
    }
    (!metadata.is_empty()).then_some(Value::Object(metadata))
}

pub(crate) fn merge_usage_request_metadata(
    base: Option<Value>,
    override_value: Option<Value>,
) -> Option<Value> {
    let mut metadata = Map::new();
    if let Some(Value::Object(base)) = base.as_ref() {
        copy_allowed_metadata_fields(base, &mut metadata);
    }
    if let Some(Value::Object(override_object)) = override_value.as_ref() {
        copy_allowed_metadata_fields(override_object, &mut metadata);
    }
    (!metadata.is_empty()).then_some(Value::Object(metadata))
}

pub(crate) fn merge_usage_request_metadata_owned(
    base: Option<Value>,
    override_value: Option<Value>,
) -> Option<Value> {
    let mut metadata = match base {
        Some(Value::Object(base)) => base,
        _ => Map::new(),
    };
    if let Some(Value::Object(override_object)) = override_value {
        move_allowed_metadata_fields(override_object, &mut metadata);
    }
    (!metadata.is_empty()).then_some(Value::Object(metadata))
}

pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
    let Value::Object(object) = value? else {
        return None;
    };

    let mut filtered = Map::new();
    move_allowed_metadata_fields(object, &mut filtered);

    (!filtered.is_empty()).then_some(Value::Object(filtered))
}

pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Option<Value> {
    let object = value.and_then(Value::as_object)?;

    let mut filtered = Map::new();
    copy_allowed_metadata_fields(object, &mut filtered);

    (!filtered.is_empty()).then_some(Value::Object(filtered))
}

pub(crate) fn attach_provider_request_body_metadata(
    metadata: Option<Value>,
    provider_request_body: Option<&Value>,
) -> Option<Value> {
    let provider_body_is_object = provider_request_body.and_then(Value::as_object).is_some();
    let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body);
    let service_tier = extract_provider_service_tier_from_body(provider_request_body);
    if !provider_body_is_object && reasoning_effort.is_none() && service_tier.is_none() {
        return metadata;
    }
    let mut object = match metadata {
        Some(Value::Object(object)) => object,
        _ => Map::new(),
    };
    if provider_body_is_object {
        object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY);
        object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY);
    }
    if let Some(reasoning_effort) = reasoning_effort {
        object.insert(
            PROVIDER_REASONING_EFFORT_METADATA_KEY.to_string(),
            Value::String(reasoning_effort),
        );
    }
    if let Some(service_tier) = service_tier {
        object.insert(
            PROVIDER_SERVICE_TIER_METADATA_KEY.to_string(),
            Value::String(service_tier),
        );
    }
    (!object.is_empty()).then_some(Value::Object(object))
}

fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
    copy_non_empty_string(source, target, "trace_id");
    copy_non_empty_string(source, target, "client_ip");
    copy_non_empty_string(source, target, "user_agent");
    copy_non_empty_string(source, target, "client_family");
    copy_bool(source, target, "client_requested_stream");
    copy_bool(source, target, UPSTREAM_IS_STREAM_KEY);
    copy_non_null_value(source, target, "client_session_affinity");
    copy_bool(source, target, "api_key_is_standalone");
    copy_non_empty_string(source, target, "request_path");
    copy_non_empty_string(source, target, "request_query_string");
    copy_non_empty_string(source, target, "request_path_and_query");
    copy_non_empty_string(source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY);
    copy_non_empty_string(source, target, PROVIDER_SERVICE_TIER_METADATA_KEY);
    copy_number(source, target, "provider_request_body_base64_bytes");
    copy_number(source, target, "provider_response_body_base64_bytes");
    copy_number(source, target, "client_response_body_base64_bytes");
    copy_non_null_value(source, target, "body_size");
    copy_number(source, target, "client_response_status_code");
    copy_non_null_value(source, target, "billing_snapshot");
    copy_non_empty_string(source, target, "billing_snapshot_schema_version");
    copy_non_empty_string(source, target, "billing_snapshot_status");
    copy_non_null_value(source, target, "settlement_snapshot");
    copy_non_empty_string(source, target, "settlement_snapshot_schema_version");
    copy_non_null_value(source, target, "billing_dimensions");
    copy_non_empty_string(source, target, "model_id");
    copy_non_empty_string(source, target, "global_model_id");
    copy_non_empty_string(source, target, "global_model_name");
    copy_non_null_value(source, target, "dimensions");
    copy_non_null_value(source, target, "billing_rule_snapshot");
    copy_non_null_value(source, target, "scheduling_audit");
    copy_non_null_value(source, target, "tls_fingerprint");
    copy_number(source, target, "rate_multiplier");
    copy_bool(source, target, "is_free_tier");
    copy_number(source, target, "input_price_per_1m");
    copy_number(source, target, "output_price_per_1m");
    copy_number(source, target, "cache_creation_price_per_1m");
    copy_number(source, target, "cache_read_price_per_1m");
    copy_number(source, target, "price_per_request");
    copy_non_null_value(source, target, "proxy");
    copy_non_null_value(source, target, "stage_timings_ms");
    copy_non_null_value(source, target, "db_timings_ms");
    sanitize_request_path_metadata_fields(target);
}

fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
    remove_non_empty_string(&mut source, target, "trace_id");
    remove_non_empty_string(&mut source, target, "client_ip");
    remove_non_empty_string(&mut source, target, "user_agent");
    remove_non_empty_string(&mut source, target, "client_family");
    remove_bool(&mut source, target, "client_requested_stream");
    remove_bool(&mut source, target, UPSTREAM_IS_STREAM_KEY);
    remove_non_null_value(&mut source, target, "client_session_affinity");
    remove_bool(&mut source, target, "api_key_is_standalone");
    remove_non_empty_string(&mut source, target, "request_path");
    remove_non_empty_string(&mut source, target, "request_query_string");
    remove_non_empty_string(&mut source, target, "request_path_and_query");
    remove_non_empty_string(&mut source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY);
    remove_non_empty_string(&mut source, target, PROVIDER_SERVICE_TIER_METADATA_KEY);
    remove_number(&mut source, target, "provider_request_body_base64_bytes");
    remove_number(&mut source, target, "provider_response_body_base64_bytes");
    remove_number(&mut source, target, "client_response_body_base64_bytes");
    remove_non_null_value(&mut source, target, "body_size");
    remove_number(&mut source, target, "client_response_status_code");
    remove_non_null_value(&mut source, target, "billing_snapshot");
    remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
    remove_non_empty_string(&mut source, target, "billing_snapshot_status");
    remove_non_null_value(&mut source, target, "settlement_snapshot");
    remove_non_empty_string(&mut source, target, "settlement_snapshot_schema_version");
    remove_non_null_value(&mut source, target, "billing_dimensions");
    remove_non_empty_string(&mut source, target, "model_id");
    remove_non_empty_string(&mut source, target, "global_model_id");
    remove_non_empty_string(&mut source, target, "global_model_name");
    remove_non_null_value(&mut source, target, "dimensions");
    remove_non_null_value(&mut source, target, "billing_rule_snapshot");
    remove_non_null_value(&mut source, target, "scheduling_audit");
    remove_non_null_value(&mut source, target, "tls_fingerprint");
    remove_number(&mut source, target, "rate_multiplier");
    remove_bool(&mut source, target, "is_free_tier");
    remove_number(&mut source, target, "input_price_per_1m");
    remove_number(&mut source, target, "output_price_per_1m");
    remove_number(&mut source, target, "cache_creation_price_per_1m");
    remove_number(&mut source, target, "cache_read_price_per_1m");
    remove_number(&mut source, target, "price_per_request");
    remove_non_null_value(&mut source, target, "proxy");
    remove_non_null_value(&mut source, target, "stage_timings_ms");
    remove_non_null_value(&mut source, target, "db_timings_ms");
    sanitize_request_path_metadata_fields(target);
}

fn sanitize_request_path_metadata_fields(target: &mut Map<String, Value>) {
    let path = target
        .get("request_path")
        .and_then(Value::as_str)
        .and_then(sanitize_request_path);
    let query = target
        .get("request_query_string")
        .and_then(Value::as_str)
        .and_then(sanitize_request_query_string);
    let path_and_query = target
        .get("request_path_and_query")
        .and_then(Value::as_str)
        .and_then(|value| sanitize_request_path_and_query(value, None))
        .or_else(|| {
            path.as_deref()
                .and_then(|path| sanitize_request_path_and_query(path, query.as_deref()))
        });

    apply_optional_string_field(target, "request_path", path.as_deref());
    apply_optional_string_field(target, "request_query_string", query.as_deref());
    apply_optional_string_field(target, "request_path_and_query", path_and_query.as_deref());
}

fn apply_optional_string_field(target: &mut Map<String, Value>, key: &str, value: Option<&str>) {
    if let Some(value) = value {
        target.insert(key.to_string(), Value::String(value.to_string()));
    } else {
        target.remove(key);
    }
}

fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source
        .get(key)
        .and_then(Value::as_str)
        .map(str::trim)
        .filter(|value| !value.is_empty())
    else {
        return;
    };
    target.insert(
        key.to_string(),
        Value::String(truncate_usage_request_metadata_string(value)),
    );
}

fn remove_non_empty_string(
    source: &mut Map<String, Value>,
    target: &mut Map<String, Value>,
    key: &str,
) {
    let Some(Value::String(value)) = source.remove(key) else {
        return;
    };
    let Some(value) = trim_and_truncate_usage_request_metadata_string_owned(value) else {
        return;
    };
    target.insert(key.to_string(), Value::String(value));
}

fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source.get(key).filter(|value| value.is_number()) else {
        return;
    };
    target.insert(key.to_string(), value.clone());
}

fn remove_number(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source.remove(key).filter(|value| value.is_number()) else {
        return;
    };
    target.insert(key.to_string(), value);
}

fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source.get(key).filter(|value| value.is_boolean()) else {
        return;
    };
    target.insert(key.to_string(), value.clone());
}

fn remove_bool(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source.remove(key).filter(|value| value.is_boolean()) else {
        return;
    };
    target.insert(key.to_string(), value);
}

fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
    let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
        return;
    };
    target.insert(
        key.to_string(),
        sanitize_usage_request_metadata_value(value),
    );
}

fn remove_non_null_value(
    source: &mut Map<String, Value>,
    target: &mut Map<String, Value>,
    key: &str,
) {
    let Some(value) = source.remove(key).filter(|value| !value.is_null()) else {
        return;
    };
    target.insert(
        key.to_string(),
        sanitize_usage_request_metadata_value_owned(value),
    );
}

fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
    match value {
        Value::String(text) => Value::String(truncate_usage_request_metadata_string(text)),
        _ if usage_request_metadata_within_limits(value) => value.clone(),
        _ => truncated_usage_request_metadata_value(value),
    }
}

fn sanitize_usage_request_metadata_value_owned(value: Value) -> Value {
    match value {
        Value::String(text) => Value::String(truncate_usage_request_metadata_string_owned(text)),
        _ if usage_request_metadata_within_limits(&value) => value,
        _ => truncated_usage_request_metadata_value(&value),
    }
}

fn truncate_usage_request_metadata_string(value: &str) -> String {
    const TRUNCATED_SUFFIX: &str = "...[truncated]";

    if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
        return value.to_string();
    }

    let target_bytes =
        MAX_USAGE_REQUEST_METADATA_STRING_BYTES.saturating_sub(TRUNCATED_SUFFIX.len());
    let mut end = 0usize;
    for (idx, ch) in value.char_indices() {
        let next = idx + ch.len_utf8();
        if next > target_bytes {
            break;
        }
        end = next;
    }

    if end == 0 {
        return TRUNCATED_SUFFIX.to_string();
    }

    format!("{}{TRUNCATED_SUFFIX}", &value[..end])
}

fn trim_and_truncate_usage_request_metadata_string_owned(value: String) -> Option<String> {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return None;
    }
    if trimmed.len() == value.len() {
        return Some(truncate_usage_request_metadata_string_owned(value));
    }
    Some(truncate_usage_request_metadata_string(trimmed))
}

fn truncate_usage_request_metadata_string_owned(value: String) -> String {
    if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
        return value;
    }
    truncate_usage_request_metadata_string(value.as_str())
}

fn truncated_usage_request_metadata_value(value: &Value) -> Value {
    json!({
        "truncated": true,
        "reason": "usage_request_metadata_limits_exceeded",
        "max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
        "max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
        "max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
        "value_kind": usage_request_metadata_value_kind(value),
    })
}

fn usage_request_metadata_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_request_metadata_value_size_hint(current));
        if depth > MAX_USAGE_REQUEST_METADATA_DEPTH
            || nodes > MAX_USAGE_REQUEST_METADATA_NODES
            || estimated_bytes > MAX_USAGE_REQUEST_METADATA_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_request_metadata_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_request_metadata_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,
    }
}

#[cfg(test)]
mod tests {
    use aether_contracts::{ExecutionPlan, RequestBody};
    use serde_json::{json, Value};
    use std::collections::BTreeMap;

    use super::{
        attach_provider_request_body_metadata, build_usage_request_metadata_seed,
        merge_usage_request_metadata, merge_usage_request_metadata_owned,
        sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
        MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH,
        MAX_USAGE_REQUEST_METADATA_NODES,
    };

    fn sample_plan() -> ExecutionPlan {
        ExecutionPlan {
            request_id: "req-1".to_string(),
            candidate_id: Some("cand-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: None,
            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,
        }
    }

    fn sample_stage_timings_metadata() -> Value {
        json!({
            "stream_candidate_slot": 1,
            "stream_provider_in_flight": 2,
            "stream_upstream_headers": 180,
            "stream_first_data": 8210
        })
    }

    fn sample_db_timings_metadata() -> Value {
        json!({
            "query_count": 2,
            "query_total": 950,
            "query_max": 650,
            "operations": {
                "request_candidate_upsert": {"count": 1, "sum": 650, "max": 650},
                "usage_upsert": {"count": 1, "sum": 300, "max": 300}
            },
            "pool": {
                "max_checked_out": 20,
                "max_pool_size": 20,
                "min_idle": 0,
                "max_connections": 20,
                "max_usage_rate": 100.0
            }
        })
    }

    #[test]
    fn sanitizes_request_metadata_to_allowlist() {
        let stage_timings_ms = sample_stage_timings_metadata();
        let db_timings_ms = sample_db_timings_metadata();
        let metadata = sanitize_usage_request_metadata(Some(json!({
            "request_id": "req-1",
            "provider_id": "provider-1",
            "provider_name": "OpenAI",
            "model": "gpt-5",
            "candidate_index": 2,
            "trace_id": "trace-1",
            "client_ip": "203.0.113.8",
            "user_agent": "Claude-Code/1.0",
            "client_requested_stream": false,
            "upstream_is_stream": true,
            "api_key_is_standalone": true,
            "provider_request_body_base64_bytes": 512,
            "provider_response_body_base64_bytes": 1024,
            "client_response_body_base64_bytes": 2048,
            "body_size": {
                "client_request_body": "1 KB",
                "provider_request_body": "4 KB",
                "provider_over_client": "4x"
            },
            "billing_snapshot": {"status": "complete"},
            "billing_snapshot_schema_version": "2.0",
            "billing_snapshot_status": "complete",
            "model_id": "model-1",
            "global_model_id": "global-model-1",
            "global_model_name": "gpt-5",
            "dimensions": {"total_input_context": 10},
            "rate_multiplier": 1.25,
            "is_free_tier": false,
            "input_price_per_1m": 3.0,
            "output_price_per_1m": 15.0,
            "cache_creation_price_per_1m": 3.75,
            "cache_read_price_per_1m": 0.3,
            "price_per_request": 0.02,
            "stage_timings_ms": stage_timings_ms.clone(),
            "db_timings_ms": db_timings_ms.clone(),
            "original_headers": {"authorization": "Bearer secret"},
            "original_request_body": {"messages": []},
            "provider_request_headers": {"authorization": "Bearer secret"},
            "upstream_url": "https://example.com/v1/chat/completions"
        })))
        .expect("metadata should remain");

        assert_eq!(
            metadata,
            json!({
                "trace_id": "trace-1",
                "client_ip": "203.0.113.8",
                "user_agent": "Claude-Code/1.0",
                "client_requested_stream": false,
                "upstream_is_stream": true,
                "api_key_is_standalone": true,
                "provider_request_body_base64_bytes": 512,
                "provider_response_body_base64_bytes": 1024,
                "client_response_body_base64_bytes": 2048,
                "body_size": {
                    "client_request_body": "1 KB",
                    "provider_request_body": "4 KB",
                    "provider_over_client": "4x"
                },
                "billing_snapshot": {"status": "complete"},
                "billing_snapshot_schema_version": "2.0",
                "billing_snapshot_status": "complete",
                "model_id": "model-1",
                "global_model_id": "global-model-1",
                "global_model_name": "gpt-5",
                "dimensions": {"total_input_context": 10},
                "rate_multiplier": 1.25,
                "is_free_tier": false,
                "input_price_per_1m": 3.0,
                "output_price_per_1m": 15.0,
                "cache_creation_price_per_1m": 3.75,
                "cache_read_price_per_1m": 0.3,
                "price_per_request": 0.02,
                "stage_timings_ms": stage_timings_ms,
                "db_timings_ms": db_timings_ms
            })
        );
    }

    #[test]
    fn sanitizes_request_path_query_metadata() {
        let metadata = sanitize_usage_request_metadata(Some(json!({
            "request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret",
            "request_query_string": "key=secret&alt=sse&pageSize=10&token=hidden",
            "request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse&pageSize=10&token=hidden",
        })))
        .expect("metadata should remain");

        assert_eq!(
            metadata,
            json!({
                "request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent",
                "request_query_string": "alt=sse&pageSize=10",
                "request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse&pageSize=10",
            })
        );
    }

    #[test]
    fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
        let metadata = sanitize_usage_request_metadata(Some(json!({
            "trace_id": "t".repeat(2_048),
            "billing_snapshot": {
                "payload": "x".repeat(32 * 1024)
            }
        })))
        .expect("metadata should remain");

        assert!(metadata
            .get("trace_id")
            .and_then(Value::as_str)
            .is_some_and(|value| value.ends_with("...[truncated]")));
        assert_eq!(
            metadata.get("billing_snapshot"),
            Some(&json!({
                "truncated": true,
                "reason": "usage_request_metadata_limits_exceeded",
                "max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
                "max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
                "max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
                "value_kind": "object",
            }))
        );
    }

    #[test]
    fn sanitizes_request_metadata_preserves_tls_fingerprint() {
        let metadata = sanitize_usage_request_metadata(Some(json!({
            "tls_fingerprint": {
                "incoming": {
                    "source": "forwarded_header",
                    "ja3": "incoming-ja3",
                    "ja4": "incoming-ja4"
                },
                "outgoing": {
                    "source": "aether_transport_config",
                    "backend": "reqwest_rustls",
                    "observed": false
                }
            },
            "untrusted_tls_fingerprint": {
                "ja3": "spoofed"
            }
        })))
        .expect("metadata should remain");

        assert_eq!(
            metadata,
            json!({
                "tls_fingerprint": {
                    "incoming": {
                        "source": "forwarded_header",
                        "ja3": "incoming-ja3",
                        "ja4": "incoming-ja4"
                    },
                    "outgoing": {
                        "source": "aether_transport_config",
                        "backend": "reqwest_rustls",
                        "observed": false
                    }
                }
            })
        );
    }

    #[test]
    fn builds_seed_from_context_and_allowlisted_metadata_only() {
        let metadata = build_usage_request_metadata_seed(
            &sample_plan(),
            Some(
                json!({
                    "request_id": "req-1",
                    "candidate_index": 0,
                    "client_requested_stream": false,
                    "upstream_is_stream": true,
                    "api_key_is_standalone": true,
                    "provider_id": "provider-1",
                    "model_id": "model-1",
                    "global_model_id": "global-model-1",
                    "global_model_name": "gpt-5",
                    "client_ip": "203.0.113.8",
                    "user_agent": "Claude-Code/1.0",
                    "billing_snapshot": {"status": "complete"},
                    "stage_timings_ms": {
                        "stream_candidate_slot": 0,
                        "stream_upstream_headers": 180,
                        "stream_first_data": 8210
                    },
                    "db_timings_ms": {
                        "query_count": 1,
                        "query_total": 42,
                        "query_max": 42,
                        "operations": {
                            "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
                        }
                    }
                })
                .as_object()
                .expect("object"),
            ),
        )
        .expect("metadata should remain");

        assert_eq!(
            metadata,
            json!({
                "client_requested_stream": false,
                "upstream_is_stream": true,
                "api_key_is_standalone": true,
                "model_id": "model-1",
                "global_model_id": "global-model-1",
                "global_model_name": "gpt-5",
                "client_ip": "203.0.113.8",
                "user_agent": "Claude-Code/1.0",
                "billing_snapshot": {"status": "complete"},
                "stage_timings_ms": {
                    "stream_candidate_slot": 0,
                    "stream_upstream_headers": 180,
                    "stream_first_data": 8210
                },
                "db_timings_ms": {
                    "query_count": 1,
                    "query_total": 42,
                    "query_max": 42,
                    "operations": {
                        "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
                    }
                }
            })
        );
    }

    #[test]
    fn merges_and_filters_request_metadata() {
        let metadata = merge_usage_request_metadata(
            Some(json!({
                "request_id": "req-1"
            })),
            Some(json!({
                "candidate_index": 0,
                "provider_name": "OpenAI"
            })),
        );

        assert_eq!(metadata, None);
    }

    #[test]
    fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() {
        let metadata = Some(json!({
            "trace_id": "trace-1",
            "provider_reasoning_effort": "high",
            "provider_service_tier": "priority"
        }));

        let updated = attach_provider_request_body_metadata(
            metadata.clone(),
            Some(&json!({
                "model": "gpt-5",
                "reasoning": { "effort": "low" },
                "service_tier": "standard"
            })),
        )
        .expect("metadata should remain");

        assert_eq!(
            updated,
            json!({
                "trace_id": "trace-1",
                "provider_reasoning_effort": "low",
                "provider_service_tier": "standard"
            })
        );

        let cleared = attach_provider_request_body_metadata(
            metadata,
            Some(&json!({
                "model": "gpt-5"
            })),
        )
        .expect("metadata should retain unrelated fields");

        assert_eq!(
            cleared,
            json!({
                "trace_id": "trace-1"
            })
        );
    }

    #[test]
    fn owned_merge_matches_filtered_merge_for_trusted_objects() {
        let base = Some(json!({
            "trace_id": "trace-1",
            "provider_request_body_base64_bytes": 128
        }));
        let override_value = Some(json!({
            "billing_snapshot_status": "complete",
            "trace_id": "trace-2"
        }));

        assert_eq!(
            merge_usage_request_metadata_owned(base.clone(), override_value.clone()),
            merge_usage_request_metadata(base, override_value)
        );
    }

    #[test]
    fn borrowed_sanitize_matches_owned_sanitize() {
        let value = json!({
            "trace_id": "trace-1",
            "billing_snapshot": {"status": "complete"},
            "provider_name": "OpenAI"
        });

        assert_eq!(
            sanitize_usage_request_metadata_ref(Some(&value)),
            sanitize_usage_request_metadata(Some(value))
        );
    }
}