Aether/apps/aether-gateway/src/execution_runtime/kiro_cache.rs

use std::collections::{BTreeMap, HashMap};
use std::fmt::Write as _;
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};

use aether_runtime_state::{DataLayerError, RuntimeState};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tracing::warn;

use crate::clock::current_unix_ms;

const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(300);
const ONE_HOUR_CACHE_TTL: Duration = Duration::from_secs(3600);
const MAX_ENTRIES: usize = 2048;
const KIRO_PROMPT_CACHE_INDEX_KEY: &str = "kiro:prompt-cache:index";
const PREFIX_LOOKBACK_WINDOW: usize = 20;
const TOKENS_PER_TOOL: u64 = 150;
const TOKENS_PER_MESSAGE: u64 = 4;
const INLINE_IMAGE_DATA_TOKEN_PLACEHOLDER: &str = "[inline-image-data]";
pub(crate) const KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD: &str = "kiro_simulated_cache_enabled";

static KIRO_PROMPT_CACHE_TRACKER: OnceLock<KiroPromptCacheTracker> = OnceLock::new();

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct KiroPromptCacheProfile {
    total_input_tokens: u64,
    min_cacheable_tokens: u64,
    breakpoints: Vec<KiroPromptCacheBreakpoint>,
    match_candidates: Vec<KiroPromptCacheCandidate>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct KiroPromptCacheBreakpoint {
    fingerprint: [u8; 32],
    cumulative_tokens: u64,
    ttl: Duration,
}

#[derive(Debug, Clone)]
struct KiroPromptCacheEntry {
    token_count: u64,
    ttl: Duration,
    expires_at: Instant,
}

#[derive(Debug, Clone, Copy, Deserialize, Serialize)]
struct KiroPromptCacheRuntimeEntry {
    token_count: u64,
    ttl_secs: u64,
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) struct KiroPromptCacheUsage {
    pub(crate) cache_creation_input_tokens: u64,
    pub(crate) cache_read_input_tokens: u64,
}

#[derive(Debug, Default)]
pub(crate) struct KiroPromptCacheTracker {
    entries: Mutex<HashMap<(String, [u8; 32]), KiroPromptCacheEntry>>,
}

#[derive(Debug)]
struct PendingBlock {
    value: Value,
    tokens: u64,
    breakpoint_ttl: Option<Duration>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct KiroPromptCacheCandidate {
    fingerprint: [u8; 32],
    cumulative_tokens: u64,
}

pub(crate) fn kiro_prompt_cache_tracker() -> &'static KiroPromptCacheTracker {
    KIRO_PROMPT_CACHE_TRACKER.get_or_init(KiroPromptCacheTracker::default)
}

pub(crate) async fn compute_kiro_prompt_cache_usage(
    runtime_state: &RuntimeState,
    credential_id: String,
    profile: &KiroPromptCacheProfile,
) -> KiroPromptCacheUsage {
    match compute_kiro_prompt_cache_usage_with_runtime_state(
        runtime_state,
        credential_id.as_str(),
        profile,
    )
    .await
    {
        Ok(usage) => usage,
        Err(err) => {
            warn!(
                event_name = "kiro_simulated_cache_runtime_state_failed",
                log_type = "event",
                error = ?err,
                "failed to update Kiro simulated cache runtime state; falling back to process-local tracker"
            );
            kiro_prompt_cache_tracker().compute_and_update(credential_id, profile)
        }
    }
}

async fn compute_kiro_prompt_cache_usage_with_runtime_state(
    runtime_state: &RuntimeState,
    credential_id: &str,
    profile: &KiroPromptCacheProfile,
) -> Result<KiroPromptCacheUsage, DataLayerError> {
    let last_breakpoint = match profile.breakpoints.last().copied() {
        Some(last_breakpoint) => last_breakpoint,
        None => return Ok(KiroPromptCacheUsage::default()),
    };

    let reversed_candidates = profile
        .match_candidates
        .iter()
        .rev()
        .copied()
        .collect::<Vec<_>>();
    let candidate_keys = reversed_candidates
        .iter()
        .map(|candidate| kiro_prompt_cache_runtime_key(credential_id, &candidate.fingerprint))
        .collect::<Vec<_>>();
    let candidate_values = runtime_state.kv_get_many(&candidate_keys).await?;
    let mut existing_entries = HashMap::<String, KiroPromptCacheRuntimeEntry>::new();
    let mut matched_tokens = 0u64;
    let mut matched_refresh: Option<(String, KiroPromptCacheRuntimeEntry)> = None;

    for ((candidate, key), value) in reversed_candidates
        .iter()
        .zip(candidate_keys.iter())
        .zip(candidate_values)
    {
        let Some(entry) = value
            .as_deref()
            .and_then(parse_kiro_prompt_cache_runtime_entry)
        else {
            continue;
        };
        existing_entries.insert(key.clone(), entry);
        if matched_tokens == 0 {
            matched_tokens = entry
                .token_count
                .min(candidate.cumulative_tokens)
                .min(profile.total_input_tokens);
            matched_refresh = Some((key.clone(), entry));
        }
    }

    if let Some((key, entry)) = matched_refresh {
        store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?;
    }

    let creation_tokens = last_breakpoint
        .cumulative_tokens
        .min(profile.total_input_tokens)
        .saturating_sub(matched_tokens);

    for breakpoint in &profile.breakpoints {
        let key = kiro_prompt_cache_runtime_key(credential_id, &breakpoint.fingerprint);
        let ttl_secs = breakpoint.ttl.as_secs().max(1);
        let entry = existing_entries
            .get(&key)
            .copied()
            .map(|existing| KiroPromptCacheRuntimeEntry {
                token_count: existing.token_count.max(breakpoint.cumulative_tokens),
                ttl_secs: existing.ttl_secs.max(ttl_secs),
            })
            .unwrap_or(KiroPromptCacheRuntimeEntry {
                token_count: breakpoint.cumulative_tokens,
                ttl_secs,
            });
        store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?;
    }

    trim_kiro_prompt_cache_runtime_state(runtime_state, MAX_ENTRIES).await;

    Ok(KiroPromptCacheUsage {
        cache_creation_input_tokens: creation_tokens,
        cache_read_input_tokens: matched_tokens,
    })
}

async fn store_kiro_prompt_cache_runtime_entry(
    runtime_state: &RuntimeState,
    key: &str,
    entry: KiroPromptCacheRuntimeEntry,
) -> Result<(), DataLayerError> {
    let ttl = Duration::from_secs(entry.ttl_secs.max(1));
    runtime_state
        .kv_set(
            key,
            encode_kiro_prompt_cache_runtime_entry(entry),
            Some(ttl),
        )
        .await?;

    let expires_at_ms = current_unix_ms().saturating_add(entry.ttl_secs.saturating_mul(1000));
    if let Err(err) = runtime_state
        .score_set(KIRO_PROMPT_CACHE_INDEX_KEY, key, expires_at_ms as f64)
        .await
    {
        warn!(
            event_name = "kiro_simulated_cache_index_update_failed",
            log_type = "event",
            cache_key = %key,
            error = ?err,
            "failed to update Kiro simulated cache index; cache entry was persisted but cleanup may lag"
        );
    }

    Ok(())
}

async fn trim_kiro_prompt_cache_runtime_state(runtime_state: &RuntimeState, max_entries: usize) {
    if let Err(err) = runtime_state
        .score_remove_by_score(KIRO_PROMPT_CACHE_INDEX_KEY, current_unix_ms() as f64)
        .await
    {
        warn!(
            event_name = "kiro_simulated_cache_index_expiry_trim_failed",
            log_type = "event",
            error = ?err,
            "failed to trim expired Kiro simulated cache index entries"
        );
        return;
    }

    let Ok(index_len) = runtime_state.score_len(KIRO_PROMPT_CACHE_INDEX_KEY).await else {
        return;
    };
    if index_len <= max_entries {
        return;
    }

    let Ok(all_members) = runtime_state
        .score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0)
        .await
    else {
        return;
    };
    let trim_count = index_len.saturating_sub(max_entries);
    if trim_count == 0 {
        return;
    }

    let trimmed_members = all_members.into_iter().take(trim_count).collect::<Vec<_>>();
    if let Err(err) = runtime_state.kv_delete_many(&trimmed_members).await {
        warn!(
            event_name = "kiro_simulated_cache_kv_trim_failed",
            log_type = "event",
            error = ?err,
            trim_count,
            "failed to delete trimmed Kiro simulated cache KV entries"
        );
    }
    if let Err(err) = runtime_state
        .score_remove_by_rank(KIRO_PROMPT_CACHE_INDEX_KEY, 0, trim_count as i64 - 1)
        .await
    {
        warn!(
            event_name = "kiro_simulated_cache_index_trim_failed",
            log_type = "event",
            error = ?err,
            trim_count,
            "failed to delete trimmed Kiro simulated cache index entries"
        );
    }
}

fn parse_kiro_prompt_cache_runtime_entry(value: &str) -> Option<KiroPromptCacheRuntimeEntry> {
    serde_json::from_str::<KiroPromptCacheRuntimeEntry>(value)
        .ok()
        .filter(|entry| entry.token_count > 0 && entry.ttl_secs > 0)
}

fn encode_kiro_prompt_cache_runtime_entry(entry: KiroPromptCacheRuntimeEntry) -> String {
    serde_json::to_string(&entry).unwrap_or_else(|_| {
        format!(
            r#"{{"token_count":{},"ttl_secs":{}}}"#,
            entry.token_count, entry.ttl_secs
        )
    })
}

fn kiro_prompt_cache_runtime_key(credential_id: &str, fingerprint: &[u8; 32]) -> String {
    let credential_hash: [u8; 32] = Sha256::digest(credential_id.as_bytes()).into();
    format!(
        "kiro:prompt-cache:{}:{}",
        hex_digest(&credential_hash),
        hex_digest(fingerprint)
    )
}

fn hex_digest(bytes: &[u8]) -> String {
    let mut output = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        let _ = write!(&mut output, "{byte:02x}");
    }
    output
}

pub(crate) fn build_kiro_prompt_cache_profile(
    request_body: &Value,
    total_input_tokens: u64,
) -> Option<KiroPromptCacheProfile> {
    let model = request_body
        .get("model")
        .and_then(Value::as_str)
        .unwrap_or_default();
    let flattened = flatten_cacheable_blocks(request_body);
    let automatic_ttl = extract_cache_ttl(request_body);
    if automatic_ttl.is_none() && flattened.iter().all(|block| block.breakpoint_ttl.is_none()) {
        return None;
    }

    let prelude = canonicalize_json(serde_json::json!({
        "model": request_body.get("model").cloned().unwrap_or(Value::Null),
        "tool_choice": request_body.get("tool_choice").cloned().unwrap_or(Value::Null),
    }));
    let mut prefix_hasher = Sha256::new();
    let prelude_bytes = serde_json::to_vec(&prelude).unwrap_or_default();
    prefix_hasher.update((prelude_bytes.len() as u64).to_be_bytes());
    prefix_hasher.update(prelude_bytes);

    let mut cumulative_tokens = 0u64;
    let mut breakpoints = Vec::new();
    let mut seen_fingerprints = std::collections::BTreeSet::<[u8; 32]>::new();
    let mut match_candidates = Vec::new();
    let last_block_index = flattened.len().saturating_sub(1);

    for (block_index, mut block) in flattened.into_iter().enumerate() {
        if block.breakpoint_ttl.is_none() && block_index == last_block_index {
            block.breakpoint_ttl = automatic_ttl;
        }
        cumulative_tokens = cumulative_tokens.saturating_add(block.tokens);
        let block_bytes = serde_json::to_vec(&block.value).unwrap_or_default();
        let block_hash: [u8; 32] = Sha256::digest(block_bytes).into();
        let mut next_prefix_hasher = prefix_hasher.clone();
        next_prefix_hasher.update(block_hash);
        let fingerprint: [u8; 32] = next_prefix_hasher.finalize().into();
        prefix_hasher = Sha256::new();
        prefix_hasher.update(fingerprint);

        if let Some(ttl) = block.breakpoint_ttl {
            push_breakpoint(
                &mut breakpoints,
                &mut seen_fingerprints,
                fingerprint,
                cumulative_tokens,
                ttl,
            );
        }
        push_match_candidate(&mut match_candidates, fingerprint, cumulative_tokens);
    }

    let min_cacheable_tokens = minimum_cacheable_tokens_for_model(model);
    let cacheable_breakpoints = breakpoints
        .into_iter()
        .filter(|breakpoint| breakpoint.cumulative_tokens >= min_cacheable_tokens)
        .collect::<Vec<_>>();
    if cacheable_breakpoints.is_empty() {
        return None;
    }
    let match_candidates = build_lookback_match_candidates(
        &match_candidates,
        &cacheable_breakpoints,
        min_cacheable_tokens,
    );
    Some(KiroPromptCacheProfile {
        total_input_tokens,
        min_cacheable_tokens,
        breakpoints: cacheable_breakpoints,
        match_candidates,
    })
}

fn build_lookback_match_candidates(
    candidates: &[KiroPromptCacheCandidate],
    breakpoints: &[KiroPromptCacheBreakpoint],
    min_cacheable_tokens: u64,
) -> Vec<KiroPromptCacheCandidate> {
    let mut out = Vec::new();
    let mut seen_fingerprints = std::collections::BTreeSet::<[u8; 32]>::new();

    for breakpoint in breakpoints {
        let Some(index) = candidates
            .iter()
            .position(|candidate| candidate.fingerprint == breakpoint.fingerprint)
        else {
            continue;
        };
        let start = index
            .saturating_add(1)
            .saturating_sub(PREFIX_LOOKBACK_WINDOW);
        for candidate in &candidates[start..=index] {
            if candidate.cumulative_tokens < min_cacheable_tokens
                || candidate.cumulative_tokens > breakpoint.cumulative_tokens
                || !seen_fingerprints.insert(candidate.fingerprint)
            {
                continue;
            }
            out.push(*candidate);
        }
    }

    out
}

pub(crate) fn kiro_simulated_cache_enabled_from_provider_config(config: Option<&Value>) -> bool {
    config
        .and_then(Value::as_object)
        .and_then(|config| config.get("kiro"))
        .and_then(Value::as_object)
        .and_then(|kiro| kiro.get("simulated_cache_enabled"))
        .and_then(Value::as_bool)
        .unwrap_or(false)
}

pub(crate) fn kiro_simulated_cache_enabled_from_report_context(
    report_context: Option<&Value>,
) -> bool {
    report_context
        .and_then(Value::as_object)
        .and_then(|context| context.get(KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD))
        .and_then(Value::as_bool)
        .unwrap_or(false)
}

pub(crate) fn billed_input_tokens(input_tokens: u64, usage: KiroPromptCacheUsage) -> u64 {
    input_tokens
        .saturating_sub(usage.cache_creation_input_tokens)
        .saturating_sub(usage.cache_read_input_tokens)
}

pub(crate) fn estimate_kiro_prompt_input_tokens(request_body: &Value) -> u64 {
    let system_tokens = request_body
        .get("system")
        .map(count_system_tokens)
        .unwrap_or(0);
    let message_tokens = request_body
        .get("messages")
        .and_then(Value::as_array)
        .map(|messages| count_messages_tokens(messages))
        .unwrap_or(0);
    let tool_tokens = request_body
        .get("tools")
        .and_then(Value::as_array)
        .map(|tools| tools.len() as u64 * TOKENS_PER_TOOL)
        .unwrap_or(0);

    (system_tokens + message_tokens + tool_tokens).max(1)
}

fn count_messages_tokens(messages: &[Value]) -> u64 {
    if messages.is_empty() {
        return 0;
    }
    let token_estimation_messages = messages
        .iter()
        .map(redact_inline_image_data_for_token_estimation)
        .collect::<Vec<_>>();
    serde_json::to_string(&token_estimation_messages)
        .map(|value| count_text_tokens(&value))
        .unwrap_or_else(|_| messages.iter().map(count_message_tokens).sum::<u64>())
        .saturating_add(messages.len() as u64 * TOKENS_PER_MESSAGE)
}

fn redact_inline_image_data_for_token_estimation(value: &Value) -> Value {
    redact_inline_image_data_value(value, false)
}

fn redact_inline_image_data_value(value: &Value, inside_image_source: bool) -> Value {
    match value {
        Value::Array(items) => Value::Array(
            items
                .iter()
                .map(|item| redact_inline_image_data_value(item, inside_image_source))
                .collect(),
        ),
        Value::Object(object) => {
            let image_block = object
                .get("type")
                .and_then(Value::as_str)
                .is_some_and(|kind| kind.eq_ignore_ascii_case("image"));
            let image_source = inside_image_source || object_has_image_media_type(object);
            let mut out = serde_json::Map::new();
            for (key, inner) in object {
                let redacted =
                    if image_source && is_inline_image_data_key(key) && inner.as_str().is_some() {
                        Value::String(INLINE_IMAGE_DATA_TOKEN_PLACEHOLDER.to_string())
                    } else {
                        let child_inside_image_source =
                            image_block && key.eq_ignore_ascii_case("source");
                        redact_inline_image_data_value(inner, child_inside_image_source)
                    };
                out.insert(key.clone(), redacted);
            }
            Value::Object(out)
        }
        other => other.clone(),
    }
}

fn object_has_image_media_type(object: &serde_json::Map<String, Value>) -> bool {
    object
        .get("media_type")
        .or_else(|| object.get("mediaType"))
        .and_then(Value::as_str)
        .is_some_and(|media_type| media_type.trim().to_ascii_lowercase().starts_with("image/"))
}

fn is_inline_image_data_key(key: &str) -> bool {
    key.eq_ignore_ascii_case("data") || key.eq_ignore_ascii_case("bytes")
}

fn push_breakpoint(
    breakpoints: &mut Vec<KiroPromptCacheBreakpoint>,
    seen_fingerprints: &mut std::collections::BTreeSet<[u8; 32]>,
    fingerprint: [u8; 32],
    cumulative_tokens: u64,
    ttl: Duration,
) {
    if seen_fingerprints.insert(fingerprint) {
        breakpoints.push(KiroPromptCacheBreakpoint {
            fingerprint,
            cumulative_tokens,
            ttl,
        });
    }
}

fn push_match_candidate(
    candidates: &mut Vec<KiroPromptCacheCandidate>,
    fingerprint: [u8; 32],
    cumulative_tokens: u64,
) {
    candidates.push(KiroPromptCacheCandidate {
        fingerprint,
        cumulative_tokens,
    });
}

fn flatten_cacheable_blocks(request_body: &Value) -> Vec<PendingBlock> {
    let mut blocks = Vec::new();
    if let Some(tools) = request_body.get("tools").and_then(Value::as_array) {
        for (tool_index, tool) in tools.iter().enumerate() {
            let breakpoint_ttl = extract_cache_ttl(tool);
            let mut normalized = tool.clone();
            strip_cache_control(&mut normalized);
            let value = canonicalize_json(serde_json::json!({
                "kind": "tool",
                "tool_index": tool_index,
                "tool": normalized,
            }));
            blocks.push(PendingBlock {
                tokens: TOKENS_PER_TOOL,
                value,
                breakpoint_ttl,
            });
        }
    }

    if let Some(system) = request_body.get("system") {
        match system {
            Value::Array(items) => {
                for (system_index, item) in items.iter().enumerate() {
                    let breakpoint_ttl = extract_cache_ttl(item);
                    let mut normalized = item.clone();
                    strip_cache_control(&mut normalized);
                    let value = canonicalize_json(serde_json::json!({
                        "kind": "system",
                        "system_index": system_index,
                        "block": normalized,
                    }));
                    blocks.push(PendingBlock {
                        tokens: count_system_block_tokens(item),
                        value,
                        breakpoint_ttl,
                    });
                }
            }
            Value::String(text) => {
                let value = canonicalize_json(serde_json::json!({
                    "kind": "system",
                    "system_index": 0,
                    "block": {"type": "text", "text": text},
                }));
                blocks.push(PendingBlock {
                    tokens: count_text_tokens(text),
                    value,
                    breakpoint_ttl: None,
                });
            }
            other => {
                let value = canonicalize_json(serde_json::json!({
                    "kind": "system",
                    "system_index": 0,
                    "block": other,
                }));
                blocks.push(PendingBlock {
                    tokens: count_system_block_tokens(other),
                    value,
                    breakpoint_ttl: None,
                });
            }
        }
    }

    if let Some(messages) = request_body.get("messages").and_then(Value::as_array) {
        for (message_index, message) in messages.iter().enumerate() {
            let role = message
                .get("role")
                .and_then(Value::as_str)
                .unwrap_or_default();
            let message_breakpoint_ttl = extract_cache_ttl(message);
            match message.get("content") {
                Some(Value::Array(items)) => {
                    let last_block_index = items.len().saturating_sub(1);
                    for (block_index, item) in items.iter().enumerate() {
                        let breakpoint_ttl =
                            extract_cache_ttl(item).or(if block_index == last_block_index {
                                message_breakpoint_ttl
                            } else {
                                None
                            });
                        let mut normalized = item.clone();
                        strip_cache_control(&mut normalized);
                        let value = canonicalize_json(serde_json::json!({
                            "kind": "message",
                            "message_index": message_index,
                            "role": role,
                            "block_index": block_index,
                            "block": normalized,
                        }));
                        blocks.push(PendingBlock {
                            tokens: count_message_content_tokens(item),
                            value,
                            breakpoint_ttl,
                        });
                    }
                }
                Some(Value::String(text)) => {
                    let value = canonicalize_json(serde_json::json!({
                        "kind": "message",
                        "message_index": message_index,
                        "role": role,
                        "block_index": 0,
                        "block": {"type": "text", "text": text},
                    }));
                    blocks.push(PendingBlock {
                        tokens: count_text_tokens(text),
                        value,
                        breakpoint_ttl: message_breakpoint_ttl,
                    });
                }
                Some(other) => {
                    let value = canonicalize_json(serde_json::json!({
                        "kind": "message",
                        "message_index": message_index,
                        "role": role,
                        "block_index": 0,
                        "block": other,
                    }));
                    blocks.push(PendingBlock {
                        tokens: count_message_content_tokens(other),
                        value,
                        breakpoint_ttl: message_breakpoint_ttl,
                    });
                }
                None => {}
            }
        }
    }
    blocks
}

fn extract_cache_ttl(value: &Value) -> Option<Duration> {
    let cache_control = value.get("cache_control")?.as_object()?;
    if !cache_control
        .get("type")
        .and_then(Value::as_str)
        .is_some_and(|value| value.trim().eq_ignore_ascii_case("ephemeral"))
    {
        return None;
    }
    Some(
        if cache_control
            .get("ttl")
            .and_then(Value::as_str)
            .is_some_and(|value| value.trim().eq_ignore_ascii_case("1h"))
        {
            ONE_HOUR_CACHE_TTL
        } else {
            DEFAULT_CACHE_TTL
        },
    )
}

fn strip_cache_control(value: &mut Value) {
    match value {
        Value::Array(items) => {
            for item in items {
                strip_cache_control(item);
            }
        }
        Value::Object(map) => {
            map.remove("cache_control");
            for item in map.values_mut() {
                strip_cache_control(item);
            }
        }
        _ => {}
    }
}

fn canonicalize_json(value: Value) -> Value {
    match value {
        Value::Array(items) => Value::Array(items.into_iter().map(canonicalize_json).collect()),
        Value::Object(map) => {
            let ordered: BTreeMap<_, _> = map
                .into_iter()
                .map(|(key, value)| (key, canonicalize_json(value)))
                .collect();
            let mut out = serde_json::Map::new();
            for (key, value) in ordered {
                out.insert(key, value);
            }
            Value::Object(out)
        }
        other => other,
    }
}

fn count_system_tokens(system: &Value) -> u64 {
    match system {
        Value::Null => 0,
        Value::String(text) => count_text_tokens(text),
        Value::Array(blocks) => blocks.iter().map(count_system_block_tokens).sum(),
        Value::Object(_) => count_system_block_tokens(system),
        _ => 0,
    }
}

fn count_system_block_tokens(block: &Value) -> u64 {
    block
        .get("text")
        .and_then(Value::as_str)
        .map(count_text_tokens)
        .unwrap_or_else(|| {
            block
                .get("thinking")
                .and_then(Value::as_str)
                .map(count_text_tokens)
                .unwrap_or_else(|| {
                    block
                        .get("content")
                        .map(count_message_content_tokens)
                        .unwrap_or(0)
                })
        })
}

fn count_message_tokens(message: &Value) -> u64 {
    let Some(object) = message.as_object() else {
        return 0;
    };
    let content = object.get("content");
    TOKENS_PER_MESSAGE
        + content
            .map(count_message_content_tokens)
            .unwrap_or_else(|| estimate_serialized_value_tokens(message))
}

fn count_message_content_tokens(value: &Value) -> u64 {
    match value {
        Value::Null => 0,
        Value::String(text) => count_text_tokens(text),
        Value::Array(items) => items.iter().map(count_message_content_tokens).sum(),
        Value::Object(object) => {
            if let Some(text) = object.get("text").and_then(Value::as_str) {
                return count_text_tokens(text);
            }
            if let Some(thinking) = object.get("thinking").and_then(Value::as_str) {
                return count_text_tokens(thinking);
            }
            if let Some(input) = object.get("input") {
                return estimate_serialized_value_tokens(input);
            }
            if let Some(content) = object.get("content") {
                return count_message_content_tokens(content);
            }
            0
        }
        _ => 0,
    }
}

fn estimate_serialized_value_tokens(value: &Value) -> u64 {
    serde_json::to_string(value)
        .map(|value| count_text_tokens(&value))
        .unwrap_or(1)
}

fn count_text_tokens(text: &str) -> u64 {
    if text.is_empty() {
        return 0;
    }

    let mut cjk_count = 0usize;
    let mut other_count = 0usize;
    for c in text.chars() {
        if c.is_whitespace() {
            continue;
        }
        if is_cjk(c) {
            cjk_count += 1;
        } else {
            other_count += 1;
        }
    }

    let tokens = (cjk_count as f64 / 1.5) + (other_count as f64 / 3.5);
    tokens.round() as u64
}

fn is_cjk(c: char) -> bool {
    matches!(
        c,
        '\u{4E00}'..='\u{9FFF}'
            | '\u{3400}'..='\u{4DBF}'
            | '\u{3040}'..='\u{309F}'
            | '\u{30A0}'..='\u{30FF}'
            | '\u{AC00}'..='\u{D7AF}'
            | '\u{1100}'..='\u{11FF}'
            | '\u{3130}'..='\u{318F}'
    )
}

fn minimum_cacheable_tokens_for_model(model: &str) -> u64 {
    let model = model.to_ascii_lowercase();
    if model.contains("opus") {
        4096
    } else if model.contains("haiku-3") || model.contains("haiku_3") {
        2048
    } else {
        1024
    }
}

impl KiroPromptCacheTracker {
    pub(crate) fn compute_and_update(
        &self,
        credential_id: String,
        profile: &KiroPromptCacheProfile,
    ) -> KiroPromptCacheUsage {
        self.compute_and_update_at(credential_id, profile, Instant::now())
    }

    fn compute_and_update_at(
        &self,
        credential_id: String,
        profile: &KiroPromptCacheProfile,
        now: Instant,
    ) -> KiroPromptCacheUsage {
        let Ok(mut entries) = self.entries.lock() else {
            return KiroPromptCacheUsage::default();
        };

        entries.retain(|_, entry| entry.expires_at > now);
        let last_breakpoint = profile.breakpoints.last().copied();
        let Some(last_breakpoint) = last_breakpoint else {
            return KiroPromptCacheUsage::default();
        };

        let mut matched_tokens = 0;
        for candidate in profile.match_candidates.iter().rev() {
            let key = (credential_id.clone(), candidate.fingerprint);
            let Some(entry) = entries.get_mut(&key) else {
                continue;
            };
            if entry.expires_at > now {
                entry.expires_at = entry.expires_at.max(now + entry.ttl);
                matched_tokens = entry
                    .token_count
                    .min(candidate.cumulative_tokens)
                    .min(profile.total_input_tokens);
                break;
            }
        }

        let creation_tokens = last_breakpoint
            .cumulative_tokens
            .min(profile.total_input_tokens)
            .saturating_sub(matched_tokens);

        for breakpoint in &profile.breakpoints {
            let key = (credential_id.clone(), breakpoint.fingerprint);
            match entries.get_mut(&key) {
                Some(existing) => {
                    existing.token_count = existing.token_count.max(breakpoint.cumulative_tokens);
                    existing.ttl = existing.ttl.max(breakpoint.ttl);
                    existing.expires_at = existing.expires_at.max(now + existing.ttl);
                }
                None => {
                    self.evict_to_capacity(&mut entries);
                    entries.insert(
                        key,
                        KiroPromptCacheEntry {
                            token_count: breakpoint.cumulative_tokens,
                            ttl: breakpoint.ttl,
                            expires_at: now + breakpoint.ttl,
                        },
                    );
                }
            }
        }

        KiroPromptCacheUsage {
            cache_creation_input_tokens: creation_tokens,
            cache_read_input_tokens: matched_tokens,
        }
    }

    fn evict_to_capacity(&self, entries: &mut HashMap<(String, [u8; 32]), KiroPromptCacheEntry>) {
        while MAX_ENTRIES > 0 && entries.len() >= MAX_ENTRIES {
            let Some(oldest_key) = entries
                .iter()
                .min_by_key(|(_, entry)| entry.expires_at)
                .map(|(key, _)| key.clone())
            else {
                break;
            };
            entries.remove(&oldest_key);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aether_runtime_state::MemoryRuntimeStateConfig;

    fn long_text(label: &str) -> String {
        format!("{} {}", label, "cacheable prompt chunk ".repeat(300))
    }

    #[test]
    fn profile_strips_cache_control_from_fingerprint() {
        let default_ttl_body = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("system"),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [{"role": "user", "content": "Perform a web search for the query: Shanghai weather"}],
            "tools": [{"type": "web_search_20250305", "name": "web_search"}]
        });
        let one_hour_body = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("system"),
                "cache_control": {"type": "ephemeral", "ttl": "1h"}
            }],
            "messages": [{"role": "user", "content": "Perform a web search for the query: Shanghai weather"}],
            "tools": [{"type": "web_search_20250305", "name": "web_search"}]
        });

        let default_profile = build_kiro_prompt_cache_profile(&default_ttl_body, 1800)
            .expect("default ttl body should build a cache profile");
        let one_hour_profile = build_kiro_prompt_cache_profile(&one_hour_body, 1800)
            .expect("one hour body should build a cache profile");

        assert_eq!(
            default_profile
                .breakpoints
                .last()
                .map(|value| value.fingerprint),
            one_hour_profile
                .breakpoints
                .last()
                .map(|value| value.fingerprint)
        );
        assert_eq!(
            default_profile.breakpoints.last().map(|value| value.ttl),
            Some(DEFAULT_CACHE_TTL)
        );
        assert_eq!(
            one_hour_profile.breakpoints.last().map(|value| value.ttl),
            Some(ONE_HOUR_CACHE_TTL)
        );
    }

    #[test]
    fn profile_reads_top_level_automatic_cache_control() {
        let request = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "cache_control": {"type": "ephemeral"},
            "messages": [{
                "role": "user",
                "content": long_text("automatic cached turn")
            }]
        });

        let profile =
            build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request))
                .expect("top-level cache_control should create an automatic cache profile");
        let tracker = KiroPromptCacheTracker::default();
        let usage = tracker.compute_and_update("cred".to_string(), &profile);

        assert_eq!(profile.breakpoints.len(), 1);
        assert!(usage.cache_creation_input_tokens > 0);
    }

    #[test]
    fn profile_does_not_create_message_end_breakpoints_from_explicit_cache_control() {
        let request = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("explicit cached system"),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [{
                "role": "user",
                "content": long_text("uncached later turn")
            }]
        });

        let profile =
            build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request))
                .expect("explicit cache_control should create a cache profile");

        assert_eq!(profile.breakpoints.len(), 1);
        assert!(profile.breakpoints[0].cumulative_tokens < profile.total_input_tokens);
    }

    #[tokio::test]
    async fn runtime_state_tracker_reads_cached_prefix_across_calls() {
        let request = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("runtime shared system"),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [{"role": "user", "content": "reuse runtime cache"}]
        });
        let profile =
            build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request))
                .expect("cacheable request should create a cache profile");
        let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());

        let first =
            compute_kiro_prompt_cache_usage(&runtime, "runtime-cred".to_string(), &profile).await;
        let second =
            compute_kiro_prompt_cache_usage(&runtime, "runtime-cred".to_string(), &profile).await;

        assert!(first.cache_creation_input_tokens > 0);
        assert_eq!(first.cache_read_input_tokens, 0);
        assert_eq!(second.cache_creation_input_tokens, 0);
        assert!(second.cache_read_input_tokens > 0);
    }

    #[tokio::test]
    async fn runtime_state_tracker_trims_oldest_entries_to_capacity() {
        let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
        let now_ms = current_unix_ms();
        let keys = [
            "kiro:prompt-cache:test-oldest".to_string(),
            "kiro:prompt-cache:test-middle".to_string(),
            "kiro:prompt-cache:test-newest".to_string(),
        ];

        for (index, key) in keys.iter().enumerate() {
            runtime
                .kv_set(
                    key,
                    encode_kiro_prompt_cache_runtime_entry(KiroPromptCacheRuntimeEntry {
                        token_count: 100 + index as u64,
                        ttl_secs: 120,
                    }),
                    Some(Duration::from_secs(120)),
                )
                .await
                .expect("cache entry should store");
            runtime
                .score_set(
                    KIRO_PROMPT_CACHE_INDEX_KEY,
                    key,
                    now_ms.saturating_add(60_000 + index as u64 * 1_000) as f64,
                )
                .await
                .expect("cache index should store");
        }

        trim_kiro_prompt_cache_runtime_state(&runtime, 2).await;

        assert_eq!(
            runtime
                .kv_get(&keys[0])
                .await
                .expect("oldest entry should read"),
            None
        );
        assert!(runtime
            .kv_get(&keys[1])
            .await
            .expect("middle entry should read")
            .is_some());
        assert!(runtime
            .kv_get(&keys[2])
            .await
            .expect("newest entry should read")
            .is_some());
        assert_eq!(
            runtime
                .score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0)
                .await
                .expect("cache index should read"),
            vec![keys[1].clone(), keys[2].clone()]
        );
    }

    #[test]
    fn tracker_refreshes_cached_prefix_ttl_on_read() {
        let base = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("shared system"),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [{"role": "user", "content": "Perform a web search for the query: Shanghai weather"}],
            "tools": [{"type": "web_search_20250305", "name": "web_search"}]
        });
        let extended = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "system": [{
                "type": "text",
                "text": long_text("shared system"),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [
                {"role": "user", "content": "Perform a web search for the query: Shanghai weather"},
                {"role": "assistant", "content": "Previous answer"},
                {"role": "user", "content": "Perform a web search for the query: Shanghai weather tomorrow"}
            ],
            "tools": [{"type": "web_search_20250305", "name": "web_search"}]
        });
        let base_profile =
            build_kiro_prompt_cache_profile(&base, 1800).expect("base should be cacheable");
        let extended_profile =
            build_kiro_prompt_cache_profile(&extended, 2200).expect("extended should be cacheable");
        let tracker = KiroPromptCacheTracker::default();
        let start = Instant::now();

        let first = tracker.compute_and_update_at("cred".to_string(), &base_profile, start);
        assert!(first.cache_creation_input_tokens > 0);
        assert_eq!(first.cache_read_input_tokens, 0);

        let hit = tracker.compute_and_update_at(
            "cred".to_string(),
            &extended_profile,
            start + Duration::from_secs(299),
        );
        assert!(hit.cache_read_input_tokens > 0);

        let refreshed = tracker.compute_and_update_at(
            "cred".to_string(),
            &base_profile,
            start + Duration::from_secs(301),
        );
        assert_eq!(refreshed.cache_creation_input_tokens, 0);
        assert!(refreshed.cache_read_input_tokens > 0);
    }

    #[test]
    fn tracker_reads_cached_prefix_when_cache_control_moves_to_new_tail() {
        let first = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": [{
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": long_text("shared first turn"),
                    "cache_control": {"type": "ephemeral"}
                }]
            }]
        });
        let second = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": [
                {
                    "role": "user",
                    "content": [{
                        "type": "text",
                        "text": long_text("shared first turn")
                    }]
                },
                {
                    "role": "assistant",
                    "content": "cached response"
                },
                {
                    "role": "user",
                    "content": [{
                        "type": "text",
                        "text": long_text("new tail turn"),
                        "cache_control": {"type": "ephemeral"}
                    }]
                }
            ]
        });
        let first_profile =
            build_kiro_prompt_cache_profile(&first, estimate_kiro_prompt_input_tokens(&first))
                .expect("first request should be cacheable");
        let second_profile =
            build_kiro_prompt_cache_profile(&second, estimate_kiro_prompt_input_tokens(&second))
                .expect("second request should be cacheable");
        let tracker = KiroPromptCacheTracker::default();
        let start = Instant::now();

        let created = tracker.compute_and_update_at("cred".to_string(), &first_profile, start);
        assert!(created.cache_creation_input_tokens > 0);
        assert_eq!(created.cache_read_input_tokens, 0);

        let hit = tracker.compute_and_update_at(
            "cred".to_string(),
            &second_profile,
            start + Duration::from_secs(60),
        );
        assert!(hit.cache_read_input_tokens > 0);
        assert!(hit.cache_creation_input_tokens > 0);
    }

    #[test]
    fn tracker_reads_cached_prefix_within_prompt_cache_lookback_window() {
        let first = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": [{
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": long_text("shared first turn"),
                    "cache_control": {"type": "ephemeral"}
                }]
            }]
        });
        let mut second_messages = vec![serde_json::json!({
            "role": "user",
            "content": [{
                "type": "text",
                "text": long_text("shared first turn")
            }]
        })];
        for index in 0..12 {
            second_messages.push(serde_json::json!({
                "role": if index % 2 == 0 { "assistant" } else { "user" },
                "content": format!("intermediate turn {index}")
            }));
        }
        second_messages.push(serde_json::json!({
            "role": "user",
            "content": [{
                "type": "text",
                "text": long_text("new tail turn"),
                "cache_control": {"type": "ephemeral"}
            }]
        }));
        let second = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": second_messages
        });
        let first_profile =
            build_kiro_prompt_cache_profile(&first, estimate_kiro_prompt_input_tokens(&first))
                .expect("first request should be cacheable");
        let second_profile =
            build_kiro_prompt_cache_profile(&second, estimate_kiro_prompt_input_tokens(&second))
                .expect("second request should be cacheable");
        let tracker = KiroPromptCacheTracker::default();
        let start = Instant::now();

        let created = tracker.compute_and_update_at("cred".to_string(), &first_profile, start);
        assert!(created.cache_creation_input_tokens > 0);
        assert_eq!(created.cache_read_input_tokens, 0);

        let hit = tracker.compute_and_update_at(
            "cred".to_string(),
            &second_profile,
            start + Duration::from_secs(60),
        );
        assert!(hit.cache_read_input_tokens > 0);
        assert!(hit.cache_creation_input_tokens > 0);
    }

    #[test]
    fn tracker_does_not_read_cached_prefix_outside_prompt_cache_lookback_window() {
        let first = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": [{
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": long_text("shared first turn"),
                    "cache_control": {"type": "ephemeral"}
                }]
            }]
        });
        let mut second_messages = vec![serde_json::json!({
            "role": "user",
            "content": [{
                "type": "text",
                "text": long_text("shared first turn")
            }]
        })];
        for index in 0..20 {
            second_messages.push(serde_json::json!({
                "role": if index % 2 == 0 { "assistant" } else { "user" },
                "content": format!("intermediate turn {index}")
            }));
        }
        second_messages.push(serde_json::json!({
            "role": "user",
            "content": [{
                "type": "text",
                "text": long_text("new tail turn"),
                "cache_control": {"type": "ephemeral"}
            }]
        }));
        let second = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": second_messages
        });
        let first_profile =
            build_kiro_prompt_cache_profile(&first, estimate_kiro_prompt_input_tokens(&first))
                .expect("first request should be cacheable");
        let second_profile =
            build_kiro_prompt_cache_profile(&second, estimate_kiro_prompt_input_tokens(&second))
                .expect("second request should be cacheable");
        let tracker = KiroPromptCacheTracker::default();
        let start = Instant::now();

        let created = tracker.compute_and_update_at("cred".to_string(), &first_profile, start);
        assert!(created.cache_creation_input_tokens > 0);
        assert_eq!(created.cache_read_input_tokens, 0);

        let miss = tracker.compute_and_update_at(
            "cred".to_string(),
            &second_profile,
            start + Duration::from_secs(60),
        );
        assert!(miss.cache_creation_input_tokens > 0);
        assert_eq!(miss.cache_read_input_tokens, 0);
    }

    #[test]
    fn profile_reads_message_level_cache_control() {
        let request = serde_json::json!({
            "model": "claude-sonnet-4.6",
            "messages": [{
                "role": "system",
                "content": long_text("message level cached system"),
                "cache_control": {"type": "ephemeral"}
            }]
        });

        let profile =
            build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request))
                .expect("message-level cache_control should create a cache profile");
        let tracker = KiroPromptCacheTracker::default();
        let usage = tracker.compute_and_update("cred".to_string(), &profile);

        assert!(usage.cache_creation_input_tokens > 0);
    }

    #[test]
    fn billed_input_tokens_subtracts_cache_usage() {
        assert_eq!(
            billed_input_tokens(
                100,
                KiroPromptCacheUsage {
                    cache_creation_input_tokens: 30,
                    cache_read_input_tokens: 40,
                },
            ),
            30
        );
        assert_eq!(
            billed_input_tokens(
                20,
                KiroPromptCacheUsage {
                    cache_creation_input_tokens: 30,
                    cache_read_input_tokens: 40,
                },
            ),
            0
        );
    }

    #[test]
    fn estimated_input_keeps_serialized_message_overhead_outside_cache() {
        let tracker = KiroPromptCacheTracker::default();
        let first = serde_json::json!({
            "model": "claude-sonnet-4-6",
            "messages": [{
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": "cacheable prompt chunk ".repeat(500),
                    "cache_control": {"type": "ephemeral"}
                }]
            }]
        });
        let second = serde_json::json!({
            "model": "claude-sonnet-4-6",
            "messages": [
                {
                    "role": "user",
                    "content": [{
                        "type": "text",
                        "text": "cacheable prompt chunk ".repeat(500),
                        "cache_control": {"type": "ephemeral"}
                    }]
                },
                {
                    "role": "assistant",
                    "content": "cached reply"
                },
                {
                    "role": "user",
                    "content": "new user turn"
                }
            ]
        });

        let first_estimated = estimate_kiro_prompt_input_tokens(&first);
        let first_profile = build_kiro_prompt_cache_profile(&first, first_estimated)
            .expect("first request should be cacheable");
        tracker.compute_and_update("cred".to_string(), &first_profile);

        let second_estimated = estimate_kiro_prompt_input_tokens(&second);
        let second_profile = build_kiro_prompt_cache_profile(&second, second_estimated)
            .expect("second request should be cacheable");
        let usage = tracker.compute_and_update("cred".to_string(), &second_profile);

        assert!(
            second_estimated
                > usage
                    .cache_creation_input_tokens
                    .saturating_add(usage.cache_read_input_tokens)
        );
        assert!(billed_input_tokens(second_estimated, usage) > 0);
    }

    #[test]
    fn estimated_input_tokens_include_message_overhead() {
        let request = serde_json::json!({
            "model": "claude-opus-4-7",
            "system": [{
                "type": "text",
                "text": "cacheable system ".repeat(400),
                "cache_control": {"type": "ephemeral"}
            }],
            "messages": [{
                "role": "user",
                "content": [{
                    "type": "text",
                    "text": "cacheable prompt ".repeat(800),
                    "cache_control": {"type": "ephemeral"}
                }]
            }]
        });

        let estimated = estimate_kiro_prompt_input_tokens(&request);
        let profile = build_kiro_prompt_cache_profile(&request, estimated)
            .expect("request should produce a cache profile");
        let tracker = KiroPromptCacheTracker::default();
        let usage = tracker.compute_and_update("cred".to_string(), &profile);

        let last_breakpoint_tokens = profile
            .breakpoints
            .last()
            .map(|breakpoint| breakpoint.cumulative_tokens)
            .expect("cache profile should have a breakpoint");

        assert!(estimated > last_breakpoint_tokens);
        assert!(billed_input_tokens(estimated, usage) > 0);
    }

    #[test]
    fn estimated_input_tokens_do_not_count_inline_image_base64_as_text() {
        let request = serde_json::json!({
            "model": "claude-sonnet-4-6",
            "messages": [{
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Please inspect this screenshot."
                    },
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/png",
                            "data": "a".repeat(200_000)
                        }
                    }
                ]
            }]
        });

        let estimated = estimate_kiro_prompt_input_tokens(&request);

        assert!(
            estimated < 1_000,
            "estimated input tokens should ignore inline image bytes, got {estimated}"
        );
    }

    #[test]
    fn kiro_simulated_cache_enabled_defaults_to_false_when_missing() {
        assert!(!kiro_simulated_cache_enabled_from_provider_config(None));
        assert!(!kiro_simulated_cache_enabled_from_provider_config(Some(
            &serde_json::json!({})
        )));
        assert!(!kiro_simulated_cache_enabled_from_provider_config(Some(
            &serde_json::json!({"kiro": {}})
        )));
        assert!(!kiro_simulated_cache_enabled_from_provider_config(Some(
            &serde_json::json!({"kiro": {"simulated_cache_enabled": false}})
        )));
    }

    #[test]
    fn kiro_simulated_cache_enabled_reads_nested_provider_config() {
        assert!(kiro_simulated_cache_enabled_from_provider_config(Some(
            &serde_json::json!({"kiro": {"simulated_cache_enabled": true}})
        )));
    }

    #[test]
    fn kiro_simulated_cache_enabled_reads_report_context_flag() {
        assert!(!kiro_simulated_cache_enabled_from_report_context(None));
        assert!(!kiro_simulated_cache_enabled_from_report_context(Some(
            &serde_json::json!({})
        )));
        assert!(kiro_simulated_cache_enabled_from_report_context(Some(
            &serde_json::json!({"kiro_simulated_cache_enabled": true})
        )));
    }
}