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

use std::collections::BTreeMap;

use crate::StandardizedUsage;

pub struct UsageMapper;

impl UsageMapper {
    pub fn map(
        raw_usage: &serde_json::Value,
        api_format: &str,
        extra_mapping: Option<&BTreeMap<String, String>>,
    ) -> StandardizedUsage {
        if !raw_usage.is_object() {
            return StandardizedUsage::new();
        }

        let mut usage = StandardizedUsage::new();
        let mut mapping = base_mapping(api_format);
        if let Some(extra_mapping) = extra_mapping {
            mapping.extend(extra_mapping.clone());
        }

        for (source_path, target_field) in mapping {
            if let Some(value) = get_nested_value(raw_usage, &source_path) {
                usage.set(&target_field, value.clone());
            }
        }

        apply_openai_cache_write_tokens(raw_usage, api_format, &mut usage);
        derive_missing_input_tokens(raw_usage, api_format, &mut usage);
        copy_explicit_total_tokens(raw_usage, api_format, &mut usage);
        usage.normalize_cache_creation_breakdown()
    }

    pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
        let family = api_family(api_format);
        let mut usage = if let Some(usage_value) = resolve_usage_value(response, family.as_str()) {
            Self::map(usage_value, api_format, None)
        } else {
            StandardizedUsage::new()
        };
        if is_openai_image_api(api_format) {
            apply_openai_image_response_dimensions(response, &mut usage);
        }
        usage
    }
}

fn apply_openai_cache_write_tokens(
    raw_usage: &serde_json::Value,
    api_format: &str,
    usage: &mut StandardizedUsage,
) {
    if api_family(api_format).as_str() != "openai" {
        return;
    }
    for details_key in ["prompt_tokens_details", "input_tokens_details"] {
        if let Some(value) = raw_usage
            .get(details_key)
            .and_then(|details| details.get("cache_write_tokens"))
        {
            usage.set("cache_creation_tokens", value.clone());
            return;
        }
    }
}

pub fn map_usage(raw_usage: &serde_json::Value, api_format: &str) -> StandardizedUsage {
    UsageMapper::map(raw_usage, api_format, None)
}

pub fn map_usage_from_response(
    response: &serde_json::Value,
    api_format: &str,
) -> StandardizedUsage {
    UsageMapper::map_from_response(response, api_format)
}

fn api_family(api_format: &str) -> String {
    api_format
        .split(':')
        .next()
        .unwrap_or_default()
        .trim()
        .to_ascii_lowercase()
}

fn api_kind(api_format: &str) -> String {
    api_format
        .split(':')
        .nth(1)
        .unwrap_or_default()
        .trim()
        .to_ascii_lowercase()
}

fn is_openai_image_api(api_format: &str) -> bool {
    api_family(api_format).as_str() == "openai" && api_kind(api_format).as_str() == "image"
}

fn apply_openai_image_response_dimensions(
    response: &serde_json::Value,
    usage: &mut StandardizedUsage,
) {
    let image_count = openai_image_response_image_count(response);
    if image_count <= 0 {
        return;
    }

    usage.request_count = image_count;
    usage
        .dimensions
        .insert("image_count".to_string(), serde_json::json!(image_count));
}

fn openai_image_response_image_count(response: &serde_json::Value) -> i64 {
    response
        .get("data")
        .and_then(serde_json::Value::as_array)
        .map(|items| items.len() as i64)
        .filter(|value| *value > 0)
        .or_else(|| image_result_count(response.get("result")))
        .unwrap_or(0)
}

fn image_result_count(value: Option<&serde_json::Value>) -> Option<i64> {
    match value? {
        serde_json::Value::Array(items) => Some(items.len() as i64).filter(|count| *count > 0),
        serde_json::Value::Object(object) if !object.is_empty() => Some(1),
        serde_json::Value::String(text) if !text.trim().is_empty() => Some(1),
        _ => None,
    }
}

fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
    let mut mapping = BTreeMap::new();
    match api_family(api_format).as_str() {
        "openai" => {
            mapping.insert("prompt_tokens".to_string(), "input_tokens".to_string());
            mapping.insert("completion_tokens".to_string(), "output_tokens".to_string());
            mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
            mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
            mapping.insert(
                "cache_creation_input_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "cache_creation.ephemeral_5m_input_tokens".to_string(),
                "cache_creation_ephemeral_5m_tokens".to_string(),
            );
            mapping.insert(
                "cache_creation.ephemeral_1h_input_tokens".to_string(),
                "cache_creation_ephemeral_1h_tokens".to_string(),
            );
            mapping.insert(
                "cache_read_input_tokens".to_string(),
                "cache_read_tokens".to_string(),
            );
            mapping.insert(
                "prompt_tokens_details.cached_tokens".to_string(),
                "cache_read_tokens".to_string(),
            );
            mapping.insert(
                "input_tokens_details.cached_tokens".to_string(),
                "cache_read_tokens".to_string(),
            );
            mapping.insert(
                "prompt_tokens_details.cached_creation_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "prompt_tokens_details.cache_write_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "input_tokens_details.cached_creation_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "input_tokens_details.cache_write_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "completion_tokens_details.reasoning_tokens".to_string(),
                "reasoning_tokens".to_string(),
            );
            mapping.insert(
                "output_tokens_details.reasoning_tokens".to_string(),
                "reasoning_tokens".to_string(),
            );
        }
        "gemini" => {
            mapping.insert("promptTokenCount".to_string(), "input_tokens".to_string());
            mapping.insert(
                "candidatesTokenCount".to_string(),
                "output_tokens".to_string(),
            );
            mapping.insert(
                "cachedContentTokenCount".to_string(),
                "cache_read_tokens".to_string(),
            );
            mapping.insert(
                "usageMetadata.promptTokenCount".to_string(),
                "input_tokens".to_string(),
            );
            mapping.insert(
                "usageMetadata.candidatesTokenCount".to_string(),
                "output_tokens".to_string(),
            );
            mapping.insert(
                "usageMetadata.cachedContentTokenCount".to_string(),
                "cache_read_tokens".to_string(),
            );
        }
        "claude" | "anthropic" => {
            mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
            mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
            mapping.insert(
                "cache_creation_input_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "cache_creation.ephemeral_5m_input_tokens".to_string(),
                "cache_creation_ephemeral_5m_tokens".to_string(),
            );
            mapping.insert(
                "cache_creation.ephemeral_1h_input_tokens".to_string(),
                "cache_creation_ephemeral_1h_tokens".to_string(),
            );
            mapping.insert(
                "cache_read_input_tokens".to_string(),
                "cache_read_tokens".to_string(),
            );
        }
        _ => {
            mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
            mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
            mapping.insert(
                "cache_creation_input_tokens".to_string(),
                "cache_creation_tokens".to_string(),
            );
            mapping.insert(
                "cache_read_input_tokens".to_string(),
                "cache_read_tokens".to_string(),
            );
        }
    }
    mapping
}

fn derive_missing_input_tokens(
    raw_usage: &serde_json::Value,
    api_format: &str,
    usage: &mut StandardizedUsage,
) {
    if usage.input_tokens > 0 || api_family(api_format).as_str() != "openai" {
        return;
    }

    let Some(total_tokens) = numeric_i64(raw_usage.get("total_tokens")) else {
        return;
    };
    let output_tokens = usage
        .output_tokens
        .max(numeric_i64(raw_usage.get("completion_tokens")).unwrap_or_default())
        .max(numeric_i64(raw_usage.get("output_tokens")).unwrap_or_default());
    let inferred_input_tokens = total_tokens.saturating_sub(output_tokens);
    if inferred_input_tokens > 0 {
        usage.input_tokens = inferred_input_tokens;
    }
}

fn copy_explicit_total_tokens(
    raw_usage: &serde_json::Value,
    api_format: &str,
    usage: &mut StandardizedUsage,
) {
    let total_tokens = match api_family(api_format).as_str() {
        "gemini" => numeric_i64(raw_usage.get("totalTokenCount")),
        _ => numeric_i64(raw_usage.get("total_tokens")),
    };
    if let Some(total_tokens) = total_tokens.filter(|value| *value > 0) {
        usage
            .dimensions
            .insert("total_tokens".to_string(), serde_json::json!(total_tokens));
    }
}

fn numeric_i64(value: Option<&serde_json::Value>) -> Option<i64> {
    value.and_then(|value| {
        value
            .as_i64()
            .or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
    })
}

fn get_nested_value<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
    let mut current = value;
    for segment in path.split('.') {
        current = current.get(segment)?;
    }
    Some(current)
}

fn resolve_usage_value<'a>(
    response: &'a serde_json::Value,
    family: &str,
) -> Option<&'a serde_json::Value> {
    match family {
        "gemini" => {
            if let Some(usage) = response.get("usageMetadata") {
                return Some(usage);
            }
            if let Some(usage) = response
                .get("candidates")
                .and_then(|value| value.get(0))
                .and_then(|value| value.get("usageMetadata"))
            {
                return Some(usage);
            }
        }
        _ => {
            if let Some(usage) = response.get("usage") {
                return Some(usage);
            }
        }
    }

    for nested_key in ["response", "message", "item"] {
        if let Some(nested) = response.get(nested_key) {
            if let Some(usage) = resolve_usage_value(nested, family) {
                return Some(usage);
            }
        }
    }

    if let Some(chunks) = response.get("chunks").and_then(serde_json::Value::as_array) {
        for chunk in chunks.iter().rev() {
            if let Some(usage) = resolve_usage_value(chunk, family) {
                return Some(usage);
            }
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::{map_usage, map_usage_from_response};

    #[test]
    fn maps_openai_usage() {
        let usage = map_usage(
            &serde_json::json!({
                "prompt_tokens": 12,
                "completion_tokens": 8,
                "prompt_tokens_details": {
                    "cached_tokens": 2,
                    "cached_creation_tokens": 1
                },
                "completion_tokens_details": { "reasoning_tokens": 3 }
            }),
            "openai:chat",
        );

        assert_eq!(usage.input_tokens, 12);
        assert_eq!(usage.output_tokens, 8);
        assert_eq!(usage.cache_creation_tokens, 1);
        assert_eq!(usage.cache_read_tokens, 2);
        assert_eq!(usage.reasoning_tokens, 3);
    }

    #[test]
    fn maps_openai_responses_usage_from_response() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "input_tokens": 14,
                    "output_tokens": 6,
                    "total_tokens": 20,
                    "input_tokens_details": {
                        "cached_tokens": 3,
                        "cached_creation_tokens": 2
                    },
                    "output_tokens_details": {
                        "reasoning_tokens": 1
                    }
                }
            }),
            "openai:responses",
        );

        assert_eq!(usage.input_tokens, 14);
        assert_eq!(usage.output_tokens, 6);
        assert_eq!(usage.cache_creation_tokens, 2);
        assert_eq!(usage.cache_read_tokens, 3);
        assert_eq!(usage.reasoning_tokens, 1);
    }

    #[test]
    fn maps_openai_responses_cache_write_tokens() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "input_tokens": 32_963,
                    "input_tokens_details": {
                        "cache_write_tokens": 512,
                        "cached_creation_tokens": 1,
                        "cached_tokens": 30_336
                    },
                    "output_tokens": 129,
                    "output_tokens_details": {
                        "reasoning_tokens": 8
                    },
                    "total_tokens": 33_092
                }
            }),
            "openai:responses",
        );

        assert_eq!(usage.input_tokens, 32_963);
        assert_eq!(usage.output_tokens, 129);
        assert_eq!(usage.cache_creation_tokens, 512);
        assert_eq!(usage.cache_read_tokens, 30_336);
        assert_eq!(usage.reasoning_tokens, 8);
    }

    #[test]
    fn maps_openai_responses_usage_with_missing_input_from_total() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "output_tokens": 899,
                    "total_tokens": 53_499,
                    "input_tokens_details": {
                        "cached_tokens": 52_600
                    }
                }
            }),
            "openai:responses",
        );

        assert_eq!(usage.input_tokens, 52_600);
        assert_eq!(usage.output_tokens, 899);
        assert_eq!(usage.cache_read_tokens, 52_600);
    }

    #[test]
    fn keeps_missing_input_derivation_scoped_to_openai() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "output_tokens": 7,
                    "total_tokens": 17,
                    "cache_read_input_tokens": 10
                }
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 0);
        assert_eq!(usage.output_tokens, 7);
        assert_eq!(usage.cache_read_tokens, 10);
    }

    #[test]
    fn maps_openai_responses_with_top_level_cache_fields() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "input_tokens": 6,
                    "output_tokens": 20,
                    "cache_creation_input_tokens": 42_262,
                    "cache_read_input_tokens": 0
                }
            }),
            "openai:chat",
        );

        assert_eq!(usage.input_tokens, 6);
        assert_eq!(usage.output_tokens, 20);
        assert_eq!(usage.cache_creation_tokens, 42_262);
        assert_eq!(usage.cache_read_tokens, 0);
    }

    #[test]
    fn maps_openai_responses_usage_from_stream_chunks() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "chunks": [
                    {
                        "type": "response.created",
                        "response": {
                            "id": "resp_123",
                            "object": "response"
                        }
                    },
                    {
                        "type": "response.completed",
                        "response": {
                            "id": "resp_123",
                            "object": "response",
                            "usage": {
                                "input_tokens": 9,
                                "output_tokens": 4,
                                "total_tokens": 13,
                                "input_tokens_details": {
                                    "cached_tokens": 5,
                                    "cached_creation_tokens": 2
                                },
                                "output_tokens_details": {
                                    "reasoning_tokens": 1
                                }
                            }
                        }
                    }
                ]
            }),
            "openai:responses",
        );

        assert_eq!(usage.input_tokens, 9);
        assert_eq!(usage.output_tokens, 4);
        assert_eq!(usage.cache_creation_tokens, 2);
        assert_eq!(usage.cache_read_tokens, 5);
        assert_eq!(usage.reasoning_tokens, 1);
    }

    #[test]
    fn maps_claude_usage() {
        let usage = map_usage(
            &serde_json::json!({
                "input_tokens": 10,
                "output_tokens": 5,
                "cache_creation_input_tokens": 4,
                "cache_read_input_tokens": 1
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 10);
        assert_eq!(usage.output_tokens, 5);
        assert_eq!(usage.cache_creation_tokens, 4);
        assert_eq!(usage.cache_read_tokens, 1);
    }

    #[test]
    fn maps_claude_usage_from_stream_chunks() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "chunks": [
                    {
                        "type": "message_start",
                        "message": {
                            "usage": {
                                "input_tokens": 5,
                                "cache_creation_input_tokens": 59_573,
                                "cache_read_input_tokens": 0,
                                "output_tokens": 0
                            }
                        }
                    },
                    {
                        "type": "message_delta",
                        "usage": {
                            "input_tokens": 5,
                            "cache_creation_input_tokens": 59_573,
                            "cache_read_input_tokens": 0,
                            "output_tokens": 162
                        }
                    }
                ]
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 5);
        assert_eq!(usage.output_tokens, 162);
        assert_eq!(usage.cache_creation_tokens, 59_573);
        assert_eq!(usage.cache_read_tokens, 0);
    }

    #[test]
    fn maps_claude_usage_from_message_start_chunk() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "chunks": [
                    {
                        "type": "message_start",
                        "message": {
                            "usage": {
                                "input_tokens": 5,
                                "cache_creation_input_tokens": 59_573,
                                "cache_read_input_tokens": 0,
                                "output_tokens": 0
                            }
                        }
                    }
                ]
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 5);
        assert_eq!(usage.output_tokens, 0);
        assert_eq!(usage.cache_creation_tokens, 59_573);
        assert_eq!(usage.cache_read_tokens, 0);
    }

    #[test]
    fn maps_claude_usage_with_large_cache_read_tokens_without_subtracting_input() {
        let usage = map_usage(
            &serde_json::json!({
                "input_tokens": 4941,
                "cache_creation_input_tokens": 687,
                "cache_read_input_tokens": 52873,
                "output_tokens": 973
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 4941);
        assert_eq!(usage.cache_creation_tokens, 687);
        assert_eq!(usage.cache_read_tokens, 52873);
        assert_eq!(usage.output_tokens, 973);
    }

    #[test]
    fn maps_claude_usage_with_cache_creation_total_and_zero_ttl_breakdown() {
        let usage = map_usage(
            &serde_json::json!({
                "cache_creation": {
                    "ephemeral_1h_input_tokens": 0,
                    "ephemeral_5m_input_tokens": 0
                },
                "cache_creation_input_tokens": 2051,
                "cache_read_input_tokens": 2051,
                "inference_geo": "inference_geo",
                "input_tokens": 2095,
                "output_tokens": 503,
                "server_tool_use": {
                    "web_fetch_requests": 2,
                    "web_search_requests": 0
                },
                "service_tier": "standard"
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 2095);
        assert_eq!(usage.cache_creation_tokens, 2051);
        assert_eq!(usage.cache_creation_ephemeral_5m_tokens, 0);
        assert_eq!(usage.cache_creation_ephemeral_1h_tokens, 0);
        assert_eq!(usage.cache_read_tokens, 2051);
        assert_eq!(usage.output_tokens, 503);
    }

    #[test]
    fn maps_claude_usage_with_ephemeral_cache_breakdown() {
        let usage = map_usage(
            &serde_json::json!({
                "input_tokens": 1,
                "output_tokens": 8,
                "cache_creation": {
                    "ephemeral_1h_input_tokens": 0,
                    "ephemeral_5m_input_tokens": 5191
                },
                "cache_creation_input_tokens": 5191,
                "cache_read_input_tokens": 97634
            }),
            "claude:messages",
        );

        assert_eq!(usage.input_tokens, 1);
        assert_eq!(usage.output_tokens, 8);
        assert_eq!(usage.cache_creation_tokens, 5191);
        assert_eq!(usage.cache_creation_ephemeral_5m_tokens, 5191);
        assert_eq!(usage.cache_creation_ephemeral_1h_tokens, 0);
        assert_eq!(usage.cache_read_tokens, 97634);
    }

    #[test]
    fn maps_gemini_usage_from_response() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usageMetadata": {
                    "promptTokenCount": 14,
                    "candidatesTokenCount": 6,
                    "cachedContentTokenCount": 2
                }
            }),
            "gemini:generate_content",
        );

        assert_eq!(usage.input_tokens, 14);
        assert_eq!(usage.output_tokens, 6);
        assert_eq!(usage.cache_read_tokens, 2);
    }

    #[test]
    fn maps_gemini_usage_from_stream_chunks() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "chunks": [
                    {
                        "candidates": [
                            {
                                "content": {
                                    "parts": [{ "text": "hello" }]
                                }
                            }
                        ]
                    },
                    {
                        "usageMetadata": {
                            "promptTokenCount": 14,
                            "candidatesTokenCount": 6,
                            "cachedContentTokenCount": 2
                        }
                    }
                ]
            }),
            "gemini:generate_content",
        );

        assert_eq!(usage.input_tokens, 14);
        assert_eq!(usage.output_tokens, 6);
        assert_eq!(usage.cache_read_tokens, 2);
    }

    #[test]
    fn maps_openai_image_response_dimensions_without_usage() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "created": 1_700_000_000,
                "data": [
                    { "b64_json": "abc" },
                    { "url": "https://example.test/image.png" }
                ]
            }),
            "openai:image",
        );

        assert_eq!(usage.request_count, 2);
        assert_eq!(
            usage.dimensions.get("image_count"),
            Some(&serde_json::json!(2))
        );
    }

    #[test]
    fn maps_openai_image_response_dimensions_with_native_usage() {
        let usage = map_usage_from_response(
            &serde_json::json!({
                "usage": {
                    "input_tokens": 11,
                    "output_tokens": 22,
                    "total_tokens": 33
                },
                "data": [{ "b64_json": "abc" }]
            }),
            "openai:image",
        );

        assert_eq!(usage.input_tokens, 11);
        assert_eq!(usage.output_tokens, 22);
        assert_eq!(usage.request_count, 1);
        assert_eq!(
            usage.dimensions.get("image_count"),
            Some(&serde_json::json!(1))
        );
    }
}