fix(kiro): 忽略图片 base64 token 估算

This commit is contained in:
Entropy.Xu
2026-05-30 02:59:04 +08:00
parent 9616f458de
commit f27485ec05
@@ -11,6 +11,7 @@ const MAX_ENTRIES: usize = 2048;
const PREFIX_LOOKBACK_LIMIT: usize = 10;
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();
@@ -194,12 +195,64 @@ fn count_messages_tokens(messages: &[Value]) -> u64 {
if messages.is_empty() {
return 0;
}
serde_json::to_string(messages)
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]>,
@@ -928,6 +981,37 @@ mod tests {
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));