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 = OnceLock::new(); #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct KiroPromptCacheProfile { total_input_tokens: u64, min_cacheable_tokens: u64, breakpoints: Vec, match_candidates: Vec, } #[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>, } #[derive(Debug)] struct PendingBlock { value: Value, tokens: u64, breakpoint_ttl: Option, } #[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 { 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::>(); let candidate_keys = reversed_candidates .iter() .map(|candidate| kiro_prompt_cache_runtime_key(credential_id, &candidate.fingerprint)) .collect::>(); let candidate_values = runtime_state.kv_get_many(&candidate_keys).await?; let mut existing_entries = HashMap::::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::>(); 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 { serde_json::from_str::(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 { 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::>(); 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 { 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::>(); serde_json::to_string(&token_estimation_messages) .map(|value| count_text_tokens(&value)) .unwrap_or_else(|_| messages.iter().map(count_message_tokens).sum::()) .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) -> 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, 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, fingerprint: [u8; 32], cumulative_tokens: u64, ) { candidates.push(KiroPromptCacheCandidate { fingerprint, cumulative_tokens, }); } fn flatten_cacheable_blocks(request_body: &Value) -> Vec { 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 { 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}) ))); } }