mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge branch 'pr-506'
This commit is contained in:
853
apps/aether-gateway/src/execution_runtime/kiro_cache.rs
Normal file
853
apps/aether-gateway/src/execution_runtime/kiro_cache.rs
Normal file
@@ -0,0 +1,853 @@
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
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 PREFIX_LOOKBACK_LIMIT: usize = 10;
|
||||
const TOKENS_PER_TOOL: u64 = 150;
|
||||
const TOKENS_PER_MESSAGE: u64 = 4;
|
||||
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>,
|
||||
}
|
||||
|
||||
#[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, 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>,
|
||||
is_message_end: bool,
|
||||
}
|
||||
|
||||
pub(crate) fn kiro_prompt_cache_tracker() -> &'static KiroPromptCacheTracker {
|
||||
KIRO_PROMPT_CACHE_TRACKER.get_or_init(KiroPromptCacheTracker::default)
|
||||
}
|
||||
|
||||
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);
|
||||
if 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 active_ttl: Option<Duration> = None;
|
||||
let mut breakpoints = Vec::new();
|
||||
let mut seen_fingerprints = std::collections::BTreeSet::<[u8; 32]>::new();
|
||||
|
||||
for block in flattened {
|
||||
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 {
|
||||
active_ttl = Some(ttl);
|
||||
push_breakpoint(
|
||||
&mut breakpoints,
|
||||
&mut seen_fingerprints,
|
||||
fingerprint,
|
||||
cumulative_tokens,
|
||||
ttl,
|
||||
);
|
||||
}
|
||||
if block.is_message_end {
|
||||
if let Some(ttl) = active_ttl {
|
||||
push_breakpoint(
|
||||
&mut breakpoints,
|
||||
&mut seen_fingerprints,
|
||||
fingerprint,
|
||||
cumulative_tokens,
|
||||
ttl,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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<_>>();
|
||||
(!cacheable_breakpoints.is_empty()).then_some(KiroPromptCacheProfile {
|
||||
total_input_tokens,
|
||||
min_cacheable_tokens,
|
||||
breakpoints: cacheable_breakpoints,
|
||||
})
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
serde_json::to_string(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 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 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,
|
||||
is_message_end: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
is_message_end: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
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,
|
||||
is_message_end: false,
|
||||
});
|
||||
}
|
||||
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,
|
||||
is_message_end: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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();
|
||||
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);
|
||||
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,
|
||||
is_message_end: block_index == last_block_index,
|
||||
});
|
||||
}
|
||||
}
|
||||
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: None,
|
||||
is_message_end: true,
|
||||
});
|
||||
}
|
||||
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: None,
|
||||
is_message_end: true,
|
||||
});
|
||||
}
|
||||
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 breakpoint in profile.breakpoints.iter().rev().take(PREFIX_LOOKBACK_LIMIT) {
|
||||
let key = (credential_id.clone(), breakpoint.fingerprint);
|
||||
let Some(entry) = entries.get(&key) else {
|
||||
continue;
|
||||
};
|
||||
if entry.expires_at > now {
|
||||
matched_tokens = entry
|
||||
.token_count
|
||||
.min(breakpoint.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);
|
||||
}
|
||||
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::*;
|
||||
|
||||
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 tracker_supports_prefix_hits_without_extending_expiry() {
|
||||
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 expired = tracker.compute_and_update_at(
|
||||
"cred".to_string(),
|
||||
&base_profile,
|
||||
start + Duration::from_secs(301),
|
||||
);
|
||||
assert!(expired.cache_creation_input_tokens > 0);
|
||||
assert_eq!(expired.cache_read_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 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})
|
||||
)));
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,11 @@ use serde_json::{json, Value};
|
||||
use tracing::{debug, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::execution_runtime::kiro_cache::{
|
||||
billed_input_tokens, build_kiro_prompt_cache_profile, estimate_kiro_prompt_input_tokens,
|
||||
kiro_prompt_cache_tracker, kiro_simulated_cache_enabled_from_provider_config,
|
||||
KiroPromptCacheProfile, KiroPromptCacheUsage,
|
||||
};
|
||||
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
|
||||
use crate::execution_runtime::transport::{
|
||||
DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
|
||||
@@ -34,6 +39,7 @@ struct KiroWebSearchRequest {
|
||||
query: String,
|
||||
model: String,
|
||||
input_tokens: u64,
|
||||
cache_profile: Option<KiroPromptCacheProfile>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -153,12 +159,25 @@ pub(crate) async fn maybe_execute_kiro_web_search_stream(
|
||||
}
|
||||
|
||||
let search_results = parse_mcp_search_results(&mcp_execution.result);
|
||||
let cache_usage = if kiro_simulated_cache_enabled(state, plan).await {
|
||||
request
|
||||
.cache_profile
|
||||
.as_ref()
|
||||
.map(|profile| {
|
||||
kiro_prompt_cache_tracker()
|
||||
.compute_and_update(kiro_cache_credential_id(plan), profile)
|
||||
})
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
KiroPromptCacheUsage::default()
|
||||
};
|
||||
let sse_body = build_web_search_sse_body(
|
||||
request.model.as_str(),
|
||||
request.query.as_str(),
|
||||
tool_use_id.as_str(),
|
||||
search_results,
|
||||
request.input_tokens,
|
||||
cache_usage,
|
||||
)
|
||||
.map_err(ExecutionRuntimeTransportError::BodyEncode)?;
|
||||
let mut synthetic_context = synthetic_report_context(report_context, mcp_execution.url);
|
||||
@@ -172,6 +191,40 @@ pub(crate) async fn maybe_execute_kiro_web_search_stream(
|
||||
}))
|
||||
}
|
||||
|
||||
async fn kiro_simulated_cache_enabled(state: &AppState, plan: &ExecutionPlan) -> bool {
|
||||
if !plan
|
||||
.provider_name
|
||||
.as_deref()
|
||||
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
match state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&plan.provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(providers) => providers
|
||||
.iter()
|
||||
.find(|provider| provider.id == plan.provider_id)
|
||||
.filter(|provider| provider.provider_type.eq_ignore_ascii_case("kiro"))
|
||||
.is_some_and(|provider| {
|
||||
kiro_simulated_cache_enabled_from_provider_config(provider.config.as_ref())
|
||||
}),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "kiro_simulated_cache_config_read_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
provider_id = %plan.provider_id,
|
||||
error = ?err,
|
||||
"failed to read Kiro simulated cache provider config; defaulting disabled"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn execute_result_body_bytes(result: &ExecutionResult) -> Vec<u8> {
|
||||
let Some(body) = result.body.as_ref() else {
|
||||
return Vec::new();
|
||||
@@ -740,10 +793,12 @@ fn detect_kiro_web_search_request(
|
||||
.or(plan.model_name.as_deref())
|
||||
.unwrap_or("claude")
|
||||
.to_string();
|
||||
let input_tokens = estimate_input_tokens(original);
|
||||
return Some(KiroWebSearchRequest {
|
||||
query,
|
||||
model,
|
||||
input_tokens: estimate_input_tokens(original),
|
||||
input_tokens,
|
||||
cache_profile: build_kiro_prompt_cache_profile(original, input_tokens),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -791,6 +846,7 @@ fn detect_kiro_web_search_from_envelope(plan: &ExecutionPlan) -> Option<KiroWebS
|
||||
query,
|
||||
model,
|
||||
input_tokens: estimate_input_tokens(body),
|
||||
cache_profile: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -823,8 +879,14 @@ fn has_only_builtin_web_search_tool(body: &Value) -> bool {
|
||||
}
|
||||
|
||||
fn extract_search_query_from_claude_request(body: &Value) -> Option<String> {
|
||||
let first_message = body.get("messages")?.as_array()?.first()?;
|
||||
let text = extract_text_content(first_message.get("content")?)?;
|
||||
let messages = body.get("messages")?.as_array()?;
|
||||
let last_user_message = messages.iter().rev().find(|message| {
|
||||
message
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|role| role.eq_ignore_ascii_case("user"))
|
||||
})?;
|
||||
let text = extract_text_content(last_user_message.get("content")?)?;
|
||||
strip_search_query_prefix(text.as_str())
|
||||
}
|
||||
|
||||
@@ -855,10 +917,7 @@ fn strip_search_query_prefix(text: &str) -> Option<String> {
|
||||
}
|
||||
|
||||
fn estimate_input_tokens(body: &Value) -> u64 {
|
||||
let size = serde_json::to_string(body)
|
||||
.map(|body| body.len() as u64)
|
||||
.unwrap_or_default();
|
||||
(size / 4).max(1)
|
||||
estimate_kiro_prompt_input_tokens(body)
|
||||
}
|
||||
|
||||
fn create_mcp_request(query: &str) -> (String, McpRequest) {
|
||||
@@ -929,14 +988,26 @@ fn execution_result_body_json(result: &ExecutionResult) -> Option<Value> {
|
||||
serde_json::from_slice(&bytes).ok()
|
||||
}
|
||||
|
||||
fn kiro_cache_credential_id(plan: &ExecutionPlan) -> String {
|
||||
format!("{}:{}:{}", plan.provider_id, plan.endpoint_id, plan.key_id)
|
||||
}
|
||||
|
||||
fn build_web_search_sse_body(
|
||||
model: &str,
|
||||
query: &str,
|
||||
tool_use_id: &str,
|
||||
search_results: Option<WebSearchResults>,
|
||||
input_tokens: u64,
|
||||
cache_usage: KiroPromptCacheUsage,
|
||||
) -> Result<Vec<u8>, serde_json::Error> {
|
||||
let events = build_web_search_events(model, query, tool_use_id, search_results, input_tokens);
|
||||
let events = build_web_search_events(
|
||||
model,
|
||||
query,
|
||||
tool_use_id,
|
||||
search_results,
|
||||
input_tokens,
|
||||
cache_usage,
|
||||
);
|
||||
crate::ai_serving::api::encode_kiro_sse_events(events)
|
||||
}
|
||||
|
||||
@@ -946,7 +1017,9 @@ fn build_web_search_events(
|
||||
tool_use_id: &str,
|
||||
search_results: Option<WebSearchResults>,
|
||||
input_tokens: u64,
|
||||
cache_usage: KiroPromptCacheUsage,
|
||||
) -> Vec<Value> {
|
||||
let billed_input = billed_input_tokens(input_tokens, cache_usage);
|
||||
let message_id = format!("msg_{}", &Uuid::new_v4().simple().to_string()[..24]);
|
||||
let mut events = vec![
|
||||
json!({
|
||||
@@ -960,10 +1033,10 @@ fn build_web_search_events(
|
||||
"stop_reason": Value::Null,
|
||||
"stop_sequence": Value::Null,
|
||||
"usage": {
|
||||
"input_tokens": input_tokens,
|
||||
"input_tokens": billed_input,
|
||||
"output_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_read_input_tokens": 0
|
||||
"cache_creation_input_tokens": cache_usage.cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": cache_usage.cache_read_input_tokens
|
||||
}
|
||||
}
|
||||
}),
|
||||
@@ -1046,8 +1119,10 @@ fn build_web_search_events(
|
||||
"stop_sequence": Value::Null
|
||||
},
|
||||
"usage": {
|
||||
"input_tokens": input_tokens,
|
||||
"input_tokens": billed_input,
|
||||
"output_tokens": estimate_text_tokens(summary.as_str()),
|
||||
"cache_creation_input_tokens": cache_usage.cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": cache_usage.cache_read_input_tokens,
|
||||
"server_tool_use": {
|
||||
"web_search_requests": 1
|
||||
}
|
||||
@@ -1178,7 +1253,7 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
build_mcp_headers_from_plan, build_web_search_sse_body, detect_kiro_web_search_request,
|
||||
parse_mcp_search_results, strip_search_query_prefix,
|
||||
parse_mcp_search_results, strip_search_query_prefix, KiroPromptCacheUsage,
|
||||
};
|
||||
|
||||
fn sample_plan(body: serde_json::Value) -> ExecutionPlan {
|
||||
@@ -1234,13 +1309,26 @@ mod tests {
|
||||
"client_api_format": "claude:messages",
|
||||
"original_request_body": {
|
||||
"model": "claude-haiku-4-5-20251001",
|
||||
"messages": [{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "Perform a web search for the query: Old query"
|
||||
}]
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Earlier answer"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "Perform a web search for the query: Shanghai weather today"
|
||||
}]
|
||||
}],
|
||||
}
|
||||
],
|
||||
"tools": [{
|
||||
"type": "web_search_20250305",
|
||||
"name": "web_search",
|
||||
@@ -1389,6 +1477,7 @@ mod tests {
|
||||
"srvtoolu_123",
|
||||
None,
|
||||
12,
|
||||
KiroPromptCacheUsage::default(),
|
||||
)
|
||||
.expect("sse should encode");
|
||||
let text = String::from_utf8(sse).expect("sse should be utf8");
|
||||
@@ -1396,4 +1485,25 @@ mod tests {
|
||||
assert!(text.contains("\"type\":\"web_search_tool_result\""));
|
||||
assert!(text.contains("\"web_search_requests\":1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn web_search_sse_includes_cache_usage_in_start_and_delta() {
|
||||
let sse = build_web_search_sse_body(
|
||||
"claude-sonnet-4.6",
|
||||
"Shanghai weather",
|
||||
"srvtoolu_123",
|
||||
None,
|
||||
30,
|
||||
KiroPromptCacheUsage {
|
||||
cache_creation_input_tokens: 7,
|
||||
cache_read_input_tokens: 9,
|
||||
},
|
||||
)
|
||||
.expect("sse should encode");
|
||||
let text = String::from_utf8(sse).expect("sse should be utf8");
|
||||
|
||||
assert_eq!(text.matches("\"cache_creation_input_tokens\":7").count(), 2);
|
||||
assert_eq!(text.matches("\"cache_read_input_tokens\":9").count(), 2);
|
||||
assert_eq!(text.matches("\"input_tokens\":14").count(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ mod chatgpt_web_image;
|
||||
mod constants;
|
||||
mod fallback;
|
||||
mod grok;
|
||||
mod kiro_cache;
|
||||
mod kiro_web_search;
|
||||
pub(crate) mod ndjson;
|
||||
mod oauth_retry;
|
||||
|
||||
@@ -7,8 +7,8 @@ use std::sync::{
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame,
|
||||
StreamFramePayload,
|
||||
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage,
|
||||
StreamFrame, StreamFramePayload,
|
||||
};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
|
||||
@@ -61,6 +61,13 @@ use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::build_direct_execution_frame_stream;
|
||||
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream;
|
||||
use crate::execution_runtime::grok::maybe_execute_grok_stream;
|
||||
use crate::execution_runtime::kiro_cache::{
|
||||
billed_input_tokens as kiro_billed_input_tokens, build_kiro_prompt_cache_profile,
|
||||
estimate_kiro_prompt_input_tokens, kiro_prompt_cache_tracker,
|
||||
kiro_simulated_cache_enabled_from_provider_config,
|
||||
kiro_simulated_cache_enabled_from_report_context, KiroPromptCacheUsage,
|
||||
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD,
|
||||
};
|
||||
use crate::execution_runtime::kiro_web_search::maybe_execute_kiro_web_search_stream;
|
||||
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
|
||||
#[cfg(test)]
|
||||
@@ -311,6 +318,268 @@ fn build_stream_usage_payload(
|
||||
}
|
||||
}
|
||||
|
||||
fn seed_kiro_report_context_input_tokens(plan: &ExecutionPlan, report_context: &mut Option<Value>) {
|
||||
if !plan
|
||||
.provider_name
|
||||
.as_deref()
|
||||
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
if context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|input_tokens| input_tokens > 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(original_request_body) = context.get("original_request_body").cloned() else {
|
||||
return;
|
||||
};
|
||||
let estimated_input_tokens = estimate_kiro_prompt_input_tokens(&original_request_body);
|
||||
context.insert(
|
||||
"input_tokens".to_string(),
|
||||
Value::from(estimated_input_tokens),
|
||||
);
|
||||
}
|
||||
|
||||
async fn seed_kiro_simulated_cache_enabled(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: &mut Option<Value>,
|
||||
) {
|
||||
if !plan
|
||||
.provider_name
|
||||
.as_deref()
|
||||
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let enabled = match state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&plan.provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(providers) => providers
|
||||
.iter()
|
||||
.find(|provider| provider.id == plan.provider_id)
|
||||
.filter(|provider| provider.provider_type.eq_ignore_ascii_case("kiro"))
|
||||
.is_some_and(|provider| {
|
||||
kiro_simulated_cache_enabled_from_provider_config(provider.config.as_ref())
|
||||
}),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "kiro_simulated_cache_config_read_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
provider_id = %plan.provider_id,
|
||||
error = ?err,
|
||||
"failed to read Kiro simulated cache provider config; defaulting disabled"
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
if enabled {
|
||||
context.insert(
|
||||
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD.to_string(),
|
||||
Value::Bool(true),
|
||||
);
|
||||
} else {
|
||||
context.remove(KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD);
|
||||
}
|
||||
}
|
||||
|
||||
fn seed_kiro_report_context_prompt_cache_usage(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: &mut Option<Value>,
|
||||
) {
|
||||
if !plan
|
||||
.provider_name
|
||||
.as_deref()
|
||||
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let simulated_cache_enabled =
|
||||
kiro_simulated_cache_enabled_from_report_context(report_context.as_ref());
|
||||
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
if context
|
||||
.get("kiro_web_search_mcp")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return;
|
||||
}
|
||||
if !simulated_cache_enabled {
|
||||
return;
|
||||
}
|
||||
if kiro_cache_usage_from_context_object(context).is_some() {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(original_request_body) = context.get("original_request_body").cloned() else {
|
||||
return;
|
||||
};
|
||||
let input_tokens = context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or_else(|| {
|
||||
let estimated = estimate_kiro_prompt_input_tokens(&original_request_body);
|
||||
context.insert("input_tokens".to_string(), Value::from(estimated));
|
||||
estimated
|
||||
});
|
||||
let Some(profile) = build_kiro_prompt_cache_profile(&original_request_body, input_tokens)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let cache_usage = kiro_prompt_cache_tracker()
|
||||
.compute_and_update(kiro_stream_cache_credential_id(plan), &profile);
|
||||
if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 {
|
||||
return;
|
||||
}
|
||||
context.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
Value::from(cache_usage.cache_creation_input_tokens),
|
||||
);
|
||||
context.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
Value::from(cache_usage.cache_read_input_tokens),
|
||||
);
|
||||
}
|
||||
|
||||
fn kiro_stream_cache_credential_id(plan: &ExecutionPlan) -> String {
|
||||
format!("{}:{}:{}", plan.provider_id, plan.endpoint_id, plan.key_id)
|
||||
}
|
||||
|
||||
fn kiro_cache_usage_from_context_object(
|
||||
context: &serde_json::Map<String, Value>,
|
||||
) -> Option<KiroPromptCacheUsage> {
|
||||
let cache_creation_input_tokens = context
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_read_input_tokens = context
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
(cache_creation_input_tokens > 0 || cache_read_input_tokens > 0).then_some(
|
||||
KiroPromptCacheUsage {
|
||||
cache_creation_input_tokens,
|
||||
cache_read_input_tokens,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn kiro_cache_usage_from_report_context(report_context: &Value) -> Option<KiroPromptCacheUsage> {
|
||||
report_context
|
||||
.as_object()
|
||||
.and_then(kiro_cache_usage_from_context_object)
|
||||
}
|
||||
|
||||
fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
summary: &mut Option<ExecutionStreamTerminalSummary>,
|
||||
) {
|
||||
if !plan
|
||||
.provider_name
|
||||
.as_deref()
|
||||
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(report_context) = report_context else {
|
||||
return;
|
||||
};
|
||||
let Some(original_request_body) = report_context.get("original_request_body") else {
|
||||
return;
|
||||
};
|
||||
let simulated_cache_enabled =
|
||||
kiro_simulated_cache_enabled_from_report_context(Some(report_context));
|
||||
|
||||
let summary = summary.get_or_insert_with(ExecutionStreamTerminalSummary::default);
|
||||
let usage = summary
|
||||
.standardized_usage
|
||||
.get_or_insert_with(StandardizedUsage::new);
|
||||
let estimated_input_tokens = report_context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or_else(|| {
|
||||
let estimated_input_tokens = estimate_kiro_prompt_input_tokens(original_request_body);
|
||||
if estimated_input_tokens > 0 {
|
||||
estimated_input_tokens
|
||||
} else {
|
||||
usage.input_tokens.max(0) as u64
|
||||
}
|
||||
});
|
||||
|
||||
if !simulated_cache_enabled {
|
||||
usage.cache_creation_tokens = 0;
|
||||
usage.cache_read_tokens = 0;
|
||||
if usage.input_tokens <= 0 {
|
||||
usage.input_tokens = estimated_input_tokens as i64;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(cache_usage) = kiro_cache_usage_from_report_context(report_context) {
|
||||
usage.input_tokens = kiro_billed_input_tokens(estimated_input_tokens, cache_usage) as i64;
|
||||
usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64;
|
||||
usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64;
|
||||
return;
|
||||
}
|
||||
|
||||
if usage.cache_creation_tokens > 0 || usage.cache_read_tokens > 0 {
|
||||
if usage.input_tokens <= 0 {
|
||||
usage.input_tokens = kiro_billed_input_tokens(
|
||||
estimated_input_tokens,
|
||||
KiroPromptCacheUsage {
|
||||
cache_creation_input_tokens: usage.cache_creation_tokens.max(0) as u64,
|
||||
cache_read_input_tokens: usage.cache_read_tokens.max(0) as u64,
|
||||
},
|
||||
) as i64;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if usage.input_tokens <= 0 {
|
||||
usage.input_tokens = estimated_input_tokens as i64;
|
||||
}
|
||||
|
||||
let Some(profile) =
|
||||
build_kiro_prompt_cache_profile(original_request_body, estimated_input_tokens)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let cache_usage = kiro_prompt_cache_tracker()
|
||||
.compute_and_update(kiro_stream_cache_credential_id(plan), &profile);
|
||||
if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let billed_input_tokens = kiro_billed_input_tokens(estimated_input_tokens, cache_usage);
|
||||
usage.input_tokens = billed_input_tokens as i64;
|
||||
usage.cache_creation_tokens = cache_usage.cache_creation_input_tokens as i64;
|
||||
usage.cache_read_tokens = cache_usage.cache_read_input_tokens as i64;
|
||||
}
|
||||
|
||||
fn append_stream_capture_bytes(
|
||||
buffer: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
@@ -480,6 +749,7 @@ pub(crate) async fn execute_execution_runtime_stream(
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let stream_started_at = Instant::now();
|
||||
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
|
||||
seed_kiro_report_context_input_tokens(&plan, &mut report_context);
|
||||
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
||||
let request_candidate_status_snapshot =
|
||||
snapshot_local_request_candidate_status(&plan, report_context.as_ref());
|
||||
@@ -1238,8 +1508,13 @@ async fn execute_stream_from_frame_stream(
|
||||
"execution runtime stream must start with headers frame".to_string(),
|
||||
));
|
||||
};
|
||||
let report_context =
|
||||
let mut report_context =
|
||||
attach_provider_response_headers_to_report_context(report_context, &headers);
|
||||
seed_kiro_report_context_input_tokens(&plan, &mut report_context);
|
||||
if status_code == 200 {
|
||||
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
|
||||
seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context);
|
||||
}
|
||||
let mut buffered_frames = VecDeque::new();
|
||||
let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None;
|
||||
if status_code == 200 && should_probe_success_failover_before_stream(&headers) {
|
||||
@@ -2918,6 +3193,12 @@ async fn execute_stream_from_frame_stream(
|
||||
return;
|
||||
}
|
||||
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&mut stream_terminal_summary,
|
||||
);
|
||||
|
||||
let should_submit_report = report_kind_owned.is_some();
|
||||
let terminal_telemetry = Some(build_terminal_stream_telemetry(
|
||||
stream_started_at_for_report,
|
||||
@@ -3081,8 +3362,9 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
build_sse_body_stream, execute_execution_runtime_stream, execute_stream_from_frame_stream,
|
||||
merge_stream_terminal_summary, should_limit_direct_finalize_prefetch,
|
||||
should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
|
||||
should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream,
|
||||
should_skip_direct_finalize_prefetch,
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
||||
@@ -3146,6 +3428,592 @@ mod tests {
|
||||
assert_eq!(merged.unknown_event_count, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cacheable system ".repeat(600),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cacheable prompt ".repeat(1200),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let report_context = json!({
|
||||
"original_request_body": request_body,
|
||||
"kiro_simulated_cache_enabled": true,
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-cache-stream".into(),
|
||||
candidate_id: Some("cand-kiro-cache-stream".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-cache-stream".into(),
|
||||
endpoint_id: "endpoint-kiro-cache-stream".into(),
|
||||
key_id: "key-kiro-cache-stream".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut first_summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 6_000,
|
||||
output_tokens: 17,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut first_summary,
|
||||
);
|
||||
let first_usage = first_summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("first usage should exist");
|
||||
assert!(first_usage.cache_creation_tokens > 0);
|
||||
assert_eq!(first_usage.cache_read_tokens, 0);
|
||||
|
||||
let mut second_summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 6_000,
|
||||
output_tokens: 19,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut second_summary,
|
||||
);
|
||||
let second_usage = second_summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("second usage should exist");
|
||||
assert!(second_usage.cache_read_tokens > 0);
|
||||
assert_eq!(second_usage.cache_creation_tokens, 0);
|
||||
assert!(second_usage.input_tokens < 6_000);
|
||||
assert_eq!(second_usage.output_tokens, 19);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_summary_seeds_input_tokens_without_cache_control() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "non cacheable system ".repeat(400)
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "non cacheable prompt ".repeat(800)
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let report_context = json!({
|
||||
"original_request_body": request_body,
|
||||
"kiro_simulated_cache_enabled": true,
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-non-cache".into(),
|
||||
candidate_id: Some("cand-kiro-non-cache".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-non-cache".into(),
|
||||
endpoint_id: "endpoint-kiro-non-cache".into(),
|
||||
key_id: "key-kiro-non-cache".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 13,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut summary,
|
||||
);
|
||||
|
||||
let usage = summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("usage should exist");
|
||||
|
||||
assert!(usage.input_tokens > 0);
|
||||
assert_eq!(usage.cache_creation_tokens, 0);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 13);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cached system ".repeat(800)
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cached prompt ".repeat(1400)
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let report_context = json!({
|
||||
"original_request_body": request_body,
|
||||
"kiro_simulated_cache_enabled": true,
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-existing-cache".into(),
|
||||
candidate_id: Some("cand-kiro-existing-cache".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-existing-cache".into(),
|
||||
endpoint_id: "endpoint-kiro-existing-cache".into(),
|
||||
key_id: "key-kiro-existing-cache".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 23,
|
||||
cache_read_tokens: 200,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut summary,
|
||||
);
|
||||
|
||||
let usage = summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("usage should exist");
|
||||
|
||||
assert!(usage.input_tokens > 0);
|
||||
assert_eq!(usage.cache_read_tokens, 200);
|
||||
assert_eq!(usage.cache_creation_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 23);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "disabled cache summary system ".repeat(800),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "disabled cache summary prompt ".repeat(1400),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let report_context = json!({
|
||||
"original_request_body": request_body,
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-summary-cache-disabled".into(),
|
||||
candidate_id: Some("cand-kiro-summary-cache-disabled".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-summary-cache-disabled".into(),
|
||||
endpoint_id: "endpoint-kiro-summary-cache-disabled".into(),
|
||||
key_id: "key-kiro-summary-cache-disabled".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 23,
|
||||
cache_creation_tokens: 500,
|
||||
cache_read_tokens: 700,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut summary,
|
||||
);
|
||||
|
||||
let usage = summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("usage should exist");
|
||||
|
||||
assert!(usage.input_tokens > 0);
|
||||
assert_eq!(usage.cache_creation_tokens, 0);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 23);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cached history ".repeat(400),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "new user turn"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let report_context = json!({
|
||||
"original_request_body": request_body,
|
||||
"input_tokens": 24_770,
|
||||
"cache_creation_input_tokens": 175,
|
||||
"cache_read_input_tokens": 24_463,
|
||||
"kiro_simulated_cache_enabled": true
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-billed-input".into(),
|
||||
candidate_id: Some("cand-kiro-billed-input".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-billed-input".into(),
|
||||
endpoint_id: "endpoint-kiro-billed-input".into(),
|
||||
key_id: "key-kiro-billed-input".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut summary = Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: Some(StandardizedUsage {
|
||||
input_tokens: 132,
|
||||
output_tokens: 167,
|
||||
cache_creation_tokens: 175,
|
||||
cache_read_tokens: 24_463,
|
||||
..StandardizedUsage::new()
|
||||
}),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
|
||||
&plan,
|
||||
Some(&report_context),
|
||||
&mut summary,
|
||||
);
|
||||
|
||||
let usage = summary
|
||||
.as_ref()
|
||||
.and_then(|summary| summary.standardized_usage.as_ref())
|
||||
.expect("usage should exist");
|
||||
|
||||
assert_eq!(usage.input_tokens, 132);
|
||||
assert_eq!(usage.cache_creation_tokens, 175);
|
||||
assert_eq!(usage.cache_read_tokens, 24_463);
|
||||
assert_eq!(usage.output_tokens, 167);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn kiro_report_context_seeds_input_tokens_from_original_request_body() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "seeded system ".repeat(600)
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "seeded prompt ".repeat(1200)
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-seed".into(),
|
||||
candidate_id: Some("cand-kiro-seed".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-seed".into(),
|
||||
endpoint_id: "endpoint-kiro-seed".into(),
|
||||
key_id: "key-kiro-seed".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
let mut report_context = Some(json!({
|
||||
"original_request_body": request_body,
|
||||
"kiro_simulated_cache_enabled": true,
|
||||
}));
|
||||
|
||||
super::seed_kiro_report_context_input_tokens(&plan, &mut report_context);
|
||||
|
||||
let input_tokens = report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("input_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.expect("kiro input tokens should be seeded");
|
||||
assert!(input_tokens > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn kiro_report_context_seeds_prompt_cache_usage_before_stream_rewrite() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cache seed system ".repeat(600),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "cache seed prompt ".repeat(1200),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-cache-seed".into(),
|
||||
candidate_id: Some("cand-kiro-cache-seed".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-cache-seed".into(),
|
||||
endpoint_id: "endpoint-kiro-cache-seed".into(),
|
||||
key_id: "key-kiro-cache-seed".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
let mut report_context = Some(json!({
|
||||
"original_request_body": request_body,
|
||||
"kiro_simulated_cache_enabled": true,
|
||||
}));
|
||||
|
||||
super::seed_kiro_report_context_input_tokens(&plan, &mut report_context);
|
||||
super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context);
|
||||
|
||||
let context = report_context.as_ref().expect("context should exist");
|
||||
assert!(context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|value| value > 0));
|
||||
assert!(context
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|value| value > 0));
|
||||
assert_eq!(
|
||||
context
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_u64),
|
||||
Some(0)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn kiro_report_context_skips_prompt_cache_usage_when_disabled() {
|
||||
let request_body = json!({
|
||||
"model": "claude-opus-4-7",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "disabled cache system ".repeat(600),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "disabled cache prompt ".repeat(1200),
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-kiro-cache-disabled".into(),
|
||||
candidate_id: Some("cand-kiro-cache-disabled".into()),
|
||||
provider_name: Some("Kiro".into()),
|
||||
provider_id: "provider-kiro-cache-disabled".into(),
|
||||
endpoint_id: "endpoint-kiro-cache-disabled".into(),
|
||||
key_id: "key-kiro-cache-disabled".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"conversationState": {}})),
|
||||
stream: true,
|
||||
client_api_format: "claude:messages".into(),
|
||||
provider_api_format: "claude:messages".into(),
|
||||
model_name: Some("claude-opus-4-7".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
let mut report_context = Some(json!({
|
||||
"original_request_body": request_body,
|
||||
}));
|
||||
|
||||
super::seed_kiro_report_context_input_tokens(&plan, &mut report_context);
|
||||
super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context);
|
||||
|
||||
let context = report_context.as_ref().expect("context should exist");
|
||||
assert!(context
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|value| value > 0));
|
||||
assert_eq!(context.get("cache_creation_input_tokens"), None);
|
||||
assert_eq!(context.get("cache_read_input_tokens"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_prefetch_for_same_format_passthrough_event_streams() {
|
||||
assert!(should_skip_direct_finalize_prefetch(
|
||||
|
||||
@@ -132,6 +132,12 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
.and_then(|cfg| cfg.get("architecture_id"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let kiro_simulated_cache_enabled = config
|
||||
.and_then(|cfg| cfg.get("kiro"))
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|cfg| cfg.get("simulated_cache_enabled"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
let billing_type = quota_snapshot
|
||||
.map(|quota| quota.billing_type.clone())
|
||||
.or_else(|| provider.billing_type.clone());
|
||||
@@ -190,6 +196,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
"endpoint_health_details": endpoint_health_details,
|
||||
"ops_configured": ops_configured,
|
||||
"ops_architecture_id": ops_architecture_id,
|
||||
"kiro_simulated_cache_enabled": kiro_simulated_cache_enabled,
|
||||
"created_at": endpoint_timestamp_or_now(provider.created_at_unix_ms, now_unix_secs),
|
||||
"updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs),
|
||||
})
|
||||
|
||||
@@ -241,6 +241,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
|
||||
"pool_advanced": {"enabled": true},
|
||||
"failover_rules": {"strategy": "ordered"},
|
||||
"chat_pii_redaction": {"enabled": true},
|
||||
"kiro": {"simulated_cache_enabled": true},
|
||||
"provider_ops": {"architecture_id": "anyrouter"}
|
||||
})),
|
||||
);
|
||||
@@ -353,6 +354,7 @@ async fn gateway_handles_admin_provider_summary_locally_with_trusted_admin_princ
|
||||
assert_eq!(payload["ops_configured"], true);
|
||||
assert_eq!(payload["ops_architecture_id"], "anyrouter");
|
||||
assert_eq!(payload["chat_pii_redaction"], json!({"enabled": true}));
|
||||
assert_eq!(payload["kiro_simulated_cache_enabled"], true);
|
||||
assert_eq!(payload["created_at"], "2024-03-21T05:46:40Z");
|
||||
assert_eq!(payload["updated_at"], "2024-03-21T05:48:20Z");
|
||||
assert_eq!(
|
||||
|
||||
@@ -9,6 +9,18 @@ pub const KIRO_MAX_THINKING_BUFFER: usize = 1024 * 1024;
|
||||
|
||||
const KIRO_QUOTE_CHARS: &str = "`\"'\\#!@$%^&*()-_=+[]{};:<>,.?/";
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct KiroStreamCacheUsage {
|
||||
pub cache_creation_input_tokens: usize,
|
||||
pub cache_read_input_tokens: usize,
|
||||
}
|
||||
|
||||
impl KiroStreamCacheUsage {
|
||||
fn has_cache_tokens(self) -> bool {
|
||||
self.cache_creation_input_tokens > 0 || self.cache_read_input_tokens > 0
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encode_kiro_sse_events(events: Vec<Value>) -> Result<Vec<u8>, serde_json::Error> {
|
||||
let mut output = Vec::new();
|
||||
for event in events {
|
||||
@@ -30,7 +42,9 @@ pub fn build_kiro_initial_sse_events(
|
||||
message_id: &str,
|
||||
model: &str,
|
||||
estimated_input_tokens: usize,
|
||||
cache_usage: Option<KiroStreamCacheUsage>,
|
||||
) -> Vec<Value> {
|
||||
let usage = build_kiro_usage_payload(estimated_input_tokens, 1, cache_usage);
|
||||
vec![json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
@@ -41,10 +55,7 @@ pub fn build_kiro_initial_sse_events(
|
||||
"model": model,
|
||||
"stop_reason": Value::Null,
|
||||
"stop_sequence": Value::Null,
|
||||
"usage": {
|
||||
"input_tokens": estimated_input_tokens as u64,
|
||||
"output_tokens": 1,
|
||||
},
|
||||
"usage": usage,
|
||||
}
|
||||
})]
|
||||
}
|
||||
@@ -63,7 +74,9 @@ pub fn build_kiro_final_message_sse_events(
|
||||
stop_reason: &str,
|
||||
input_tokens: usize,
|
||||
output_tokens: usize,
|
||||
cache_usage: Option<KiroStreamCacheUsage>,
|
||||
) -> Vec<Value> {
|
||||
let usage = build_kiro_usage_payload(input_tokens, output_tokens, cache_usage);
|
||||
vec![
|
||||
json!({
|
||||
"type": "message_delta",
|
||||
@@ -71,15 +84,37 @@ pub fn build_kiro_final_message_sse_events(
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": Value::Null,
|
||||
},
|
||||
"usage": {
|
||||
"input_tokens": input_tokens as u64,
|
||||
"output_tokens": output_tokens as u64,
|
||||
}
|
||||
"usage": usage
|
||||
}),
|
||||
json!({"type": "message_stop"}),
|
||||
]
|
||||
}
|
||||
|
||||
fn build_kiro_usage_payload(
|
||||
input_tokens: usize,
|
||||
output_tokens: usize,
|
||||
cache_usage: Option<KiroStreamCacheUsage>,
|
||||
) -> Value {
|
||||
let billed_input_tokens = cache_usage
|
||||
.filter(|usage| usage.has_cache_tokens())
|
||||
.map(|usage| {
|
||||
input_tokens
|
||||
.saturating_sub(usage.cache_creation_input_tokens)
|
||||
.saturating_sub(usage.cache_read_input_tokens)
|
||||
})
|
||||
.unwrap_or(input_tokens);
|
||||
let mut usage = json!({
|
||||
"input_tokens": billed_input_tokens as u64,
|
||||
"output_tokens": output_tokens as u64,
|
||||
});
|
||||
if let Some(cache_usage) = cache_usage.filter(|usage| usage.has_cache_tokens()) {
|
||||
usage["cache_creation_input_tokens"] =
|
||||
json!(cache_usage.cache_creation_input_tokens as u64);
|
||||
usage["cache_read_input_tokens"] = json!(cache_usage.cache_read_input_tokens as u64);
|
||||
}
|
||||
usage
|
||||
}
|
||||
|
||||
pub fn calculate_kiro_context_input_tokens(percentage: f64) -> usize {
|
||||
((percentage * KIRO_CONTEXT_WINDOW_TOKENS) / 100.0) as usize
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ struct KiroClaudeStreamState {
|
||||
model: String,
|
||||
thinking_enabled: bool,
|
||||
estimated_input_tokens: usize,
|
||||
cache_usage: Option<super::KiroStreamCacheUsage>,
|
||||
message_id: String,
|
||||
output_tokens: usize,
|
||||
context_input_tokens: Option<usize>,
|
||||
|
||||
@@ -83,13 +83,16 @@ impl KiroClaudeStreamState {
|
||||
}
|
||||
.to_string()
|
||||
});
|
||||
let input_tokens = self
|
||||
.context_input_tokens
|
||||
.unwrap_or(self.estimated_input_tokens) as u64;
|
||||
let input_tokens = if self.estimated_input_tokens > 0 {
|
||||
self.estimated_input_tokens
|
||||
} else {
|
||||
self.context_input_tokens.unwrap_or_default()
|
||||
};
|
||||
encode_kiro_sse_events(build_kiro_final_message_sse_events(
|
||||
&stop_reason,
|
||||
input_tokens as usize,
|
||||
input_tokens,
|
||||
self.output_tokens,
|
||||
self.cache_usage,
|
||||
))
|
||||
.map_err(AiSurfaceFinalizeError::from)
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use crate::formats::shared::model_directives::model_directive_display_model_from
|
||||
use crate::formats::shared::AiSurfaceFinalizeError;
|
||||
use crate::provider_compat::kiro_stream::{
|
||||
build_kiro_initial_sse_events, build_kiro_stream_error_sse_events, encode_kiro_sse_events,
|
||||
KiroStreamCacheUsage,
|
||||
};
|
||||
|
||||
use super::super::{EventStreamDecoder, KiroClaudeStreamState, KiroToClaudeCliStreamState};
|
||||
@@ -97,10 +98,26 @@ impl KiroClaudeStreamState {
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_input_tokens = report_context
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let cache_read_input_tokens = report_context
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.map(|value| value as usize)
|
||||
.unwrap_or(0);
|
||||
let cache_usage = (cache_creation_input_tokens > 0 || cache_read_input_tokens > 0)
|
||||
.then_some(KiroStreamCacheUsage {
|
||||
cache_creation_input_tokens,
|
||||
cache_read_input_tokens,
|
||||
});
|
||||
Self {
|
||||
model,
|
||||
thinking_enabled,
|
||||
estimated_input_tokens,
|
||||
cache_usage,
|
||||
message_id: format!("msg_{}", Uuid::new_v4().simple()),
|
||||
..Self::default()
|
||||
}
|
||||
@@ -111,6 +128,7 @@ impl KiroClaudeStreamState {
|
||||
&self.message_id,
|
||||
&self.model,
|
||||
self.estimated_input_tokens,
|
||||
self.cache_usage,
|
||||
);
|
||||
let mut events = events;
|
||||
if !self.thinking_enabled {
|
||||
|
||||
@@ -112,6 +112,84 @@ fn kiro_stream_rewriter_restores_model_directive_display_model() {
|
||||
assert!(!text.contains("\"model\":\"claude-sonnet-4.5\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_emits_cache_usage_from_report_context() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"envelope_name": "kiro:generateAssistantResponse",
|
||||
"mapped_model": "claude-sonnet-4.5",
|
||||
"input_tokens": 100,
|
||||
"cache_creation_input_tokens": 25,
|
||||
"cache_read_input_tokens": 40
|
||||
});
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let first = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
&encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "Hello"}),
|
||||
),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
|
||||
assert_eq!(text.matches("\"input_tokens\":35").count(), 2);
|
||||
assert_eq!(
|
||||
text.matches("\"cache_creation_input_tokens\":25").count(),
|
||||
2
|
||||
);
|
||||
assert_eq!(text.matches("\"cache_read_input_tokens\":40").count(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_keeps_estimated_input_when_context_usage_is_cache_only() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"envelope_name": "kiro:generateAssistantResponse",
|
||||
"mapped_model": "claude-sonnet-4.5",
|
||||
"input_tokens": 24_344,
|
||||
"cache_creation_input_tokens": 293,
|
||||
"cache_read_input_tokens": 23_935
|
||||
});
|
||||
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
|
||||
let chunk = [
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("assistantResponseEvent"),
|
||||
&json!({"content": "Hello"}),
|
||||
),
|
||||
encode_event_frame(
|
||||
"event",
|
||||
Some("contextUsageEvent"),
|
||||
&json!({"contextUsagePercentage": 12.114}),
|
||||
),
|
||||
]
|
||||
.concat();
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(&report_context, &chunk)
|
||||
.expect("rewrite should succeed");
|
||||
let rest = rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed");
|
||||
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
|
||||
|
||||
assert_eq!(text.matches("\"input_tokens\":116").count(), 2);
|
||||
assert!(!text.contains("\"input_tokens\":0"));
|
||||
assert_eq!(
|
||||
text.matches("\"cache_creation_input_tokens\":293").count(),
|
||||
2
|
||||
);
|
||||
assert_eq!(text.matches("\"cache_read_input_tokens\":23935").count(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_stream_rewriter_converts_tool_use_to_claude_events() {
|
||||
let report_context = kiro_report_context(false);
|
||||
|
||||
@@ -24,6 +24,8 @@ INSERT INTO "usage" (
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
input_output_total_tokens,
|
||||
input_context_tokens,
|
||||
cache_creation_input_tokens,
|
||||
cache_creation_input_tokens_5m,
|
||||
cache_creation_input_tokens_1h,
|
||||
@@ -87,7 +89,28 @@ INSERT INTO "usage" (
|
||||
),
|
||||
COALESCE($22, 0),
|
||||
COALESCE($23, 0),
|
||||
COALESCE($24, COALESCE($22, 0) + COALESCE($23, 0)),
|
||||
COALESCE(
|
||||
$24,
|
||||
COALESCE($22, 0)
|
||||
+ COALESCE($23, 0)
|
||||
+ COALESCE(
|
||||
NULLIF(COALESCE($25, 0), 0),
|
||||
COALESCE($26, 0) + COALESCE($27, 0),
|
||||
0
|
||||
)
|
||||
+ COALESCE($28, 0)
|
||||
),
|
||||
COALESCE($22, 0) + COALESCE($23, 0),
|
||||
COALESCE(
|
||||
COALESCE($22, 0)
|
||||
+ COALESCE(
|
||||
NULLIF(COALESCE($25, 0), 0),
|
||||
COALESCE($26, 0) + COALESCE($27, 0),
|
||||
0
|
||||
)
|
||||
+ COALESCE($28, 0),
|
||||
0
|
||||
),
|
||||
COALESCE($25, 0),
|
||||
COALESCE($26, 0),
|
||||
COALESCE($27, 0),
|
||||
@@ -148,6 +171,8 @@ DO UPDATE SET
|
||||
input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_tokens, EXCLUDED.input_tokens) ELSE "usage".input_tokens END,
|
||||
output_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".output_tokens, EXCLUDED.output_tokens) ELSE "usage".output_tokens END,
|
||||
total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".total_tokens, EXCLUDED.total_tokens) ELSE "usage".total_tokens END,
|
||||
input_output_total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_output_total_tokens, EXCLUDED.input_output_total_tokens) ELSE "usage".input_output_total_tokens END,
|
||||
input_context_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_context_tokens, EXCLUDED.input_context_tokens) ELSE "usage".input_context_tokens END,
|
||||
cache_creation_input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE "usage".cache_creation_input_tokens END,
|
||||
cache_creation_input_tokens_5m = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE "usage".cache_creation_input_tokens_5m END,
|
||||
cache_creation_input_tokens_1h = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE "usage".cache_creation_input_tokens_1h END,
|
||||
|
||||
@@ -706,6 +706,8 @@ fn usage_sql_upsert_returning_includes_routing_placeholders() {
|
||||
super::UPSERT_SQL.contains("NULL::varchar AS settlement_billing_snapshot_schema_version")
|
||||
);
|
||||
assert!(super::UPSERT_SQL.contains("NULL::double precision AS settlement_input_price_per_1m"));
|
||||
assert!(super::UPSERT_SQL.contains("input_output_total_tokens"));
|
||||
assert!(super::UPSERT_SQL.contains("input_context_tokens"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -737,6 +739,8 @@ fn usage_sql_updates_usage_mirror_columns_from_terminal_events_only() {
|
||||
"input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_tokens, EXCLUDED.input_tokens) ELSE \"usage\".input_tokens END",
|
||||
"output_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".output_tokens, EXCLUDED.output_tokens) ELSE \"usage\".output_tokens END",
|
||||
"total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_tokens, EXCLUDED.total_tokens) ELSE \"usage\".total_tokens END",
|
||||
"input_output_total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_output_total_tokens, EXCLUDED.input_output_total_tokens) ELSE \"usage\".input_output_total_tokens END",
|
||||
"input_context_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_context_tokens, EXCLUDED.input_context_tokens) ELSE \"usage\".input_context_tokens END",
|
||||
"cache_creation_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE \"usage\".cache_creation_input_tokens END",
|
||||
"cache_creation_input_tokens_5m = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE \"usage\".cache_creation_input_tokens_5m END",
|
||||
"cache_creation_input_tokens_1h = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE \"usage\".cache_creation_input_tokens_1h END",
|
||||
|
||||
@@ -48,6 +48,7 @@ function normalizeProviderSummary(
|
||||
...provider,
|
||||
chat_pii_redaction: normalizeChatPiiRedactionProvider(provider.chat_pii_redaction),
|
||||
pool_advanced: normalizePoolAdvanced(provider.pool_advanced),
|
||||
kiro_simulated_cache_enabled: provider.kiro_simulated_cache_enabled ?? false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -684,6 +684,7 @@ export interface ProviderWithEndpointsSummary {
|
||||
failover_rules?: FailoverRulesConfig | null
|
||||
ops_configured: boolean // 是否配置了扩展操作(余额监控等)
|
||||
ops_architecture_id?: string // 扩展操作使用的架构 ID(如 cubence, anyrouter)
|
||||
kiro_simulated_cache_enabled?: boolean
|
||||
created_at: string
|
||||
updated_at: string
|
||||
}
|
||||
|
||||
@@ -275,6 +275,22 @@
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div
|
||||
v-if="form.provider_type === 'kiro'"
|
||||
class="flex items-center justify-between p-3 border rounded-lg bg-muted/50"
|
||||
>
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">模拟缓存模式</span>
|
||||
<p class="text-xs text-muted-foreground leading-relaxed">
|
||||
启用后仅对 Kiro 请求模拟 prompt cache 读写计量。
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="form.kiro_simulated_cache_enabled"
|
||||
@update:model-value="(v: boolean) => form.kiro_simulated_cache_enabled = v"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div class="flex items-center justify-between gap-4 p-3 border rounded-lg bg-muted/50">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">敏感信息保护</span>
|
||||
@@ -383,6 +399,8 @@ const form = ref({
|
||||
request_timeout: undefined as number | undefined,
|
||||
// 号池模式
|
||||
pool_mode_enabled: false,
|
||||
// Kiro 专属配置
|
||||
kiro_simulated_cache_enabled: false,
|
||||
})
|
||||
|
||||
// 重置表单
|
||||
@@ -409,6 +427,8 @@ function resetForm() {
|
||||
request_timeout: undefined,
|
||||
// 号池模式
|
||||
pool_mode_enabled: false,
|
||||
// Kiro 专属配置
|
||||
kiro_simulated_cache_enabled: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -439,6 +459,8 @@ function loadProviderData() {
|
||||
request_timeout: props.provider.request_timeout ?? undefined,
|
||||
// 号池模式
|
||||
pool_mode_enabled: poolAdvanced !== null,
|
||||
// Kiro 专属配置
|
||||
kiro_simulated_cache_enabled: props.provider.kiro_simulated_cache_enabled ?? false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -457,6 +479,9 @@ watch(() => form.value.provider_type, () => {
|
||||
if (!isEditMode.value) {
|
||||
form.value.pool_mode_enabled = false
|
||||
}
|
||||
if (form.value.provider_type !== 'kiro') {
|
||||
form.value.kiro_simulated_cache_enabled = false
|
||||
}
|
||||
})
|
||||
|
||||
// 提交表单
|
||||
@@ -506,6 +531,15 @@ const handleSubmit = async () => {
|
||||
pool_advanced: form.value.pool_mode_enabled
|
||||
? (currentPoolAdvanced ?? {})
|
||||
: null,
|
||||
...(form.value.provider_type === 'kiro'
|
||||
? {
|
||||
config: {
|
||||
kiro: {
|
||||
simulated_cache_enabled: form.value.kiro_simulated_cache_enabled,
|
||||
},
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
}
|
||||
|
||||
if (isEditMode.value && props.provider) {
|
||||
|
||||
Reference in New Issue
Block a user