i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-13 14:40:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

修复：完善 Kiro 模拟缓存共享回收

Browse Source
This commit is contained in:
Entropy.Xu
2026-05-31 22:43:27 +08:00
parent c2bcfab7d4
commit d1b64b6748
2 changed files with 212 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/execution_runtime/kiro_cache.rs
+154 -14
View File
@@ -9,9 +9,12 @@ use serde_json::Value;
use sha2::{Digest, Sha256};
use tracing::warn;
use crate::clock::current_unix_ms;
const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(300);
const ONE_HOUR_CACHE_TTL: Duration = Duration::from_secs(3600);
const MAX_ENTRIES: usize = 2048;
const KIRO_PROMPT_CACHE_INDEX_KEY: &str = "kiro:prompt-cache:index";
const PREFIX_LOOKBACK_WINDOW: usize = 20;
const TOKENS_PER_TOOL: u64 = 150;
const TOKENS_PER_MESSAGE: u64 = 4;
@@ -148,13 +151,7 @@ async fn compute_kiro_prompt_cache_usage_with_runtime_state(
}
if let Some((key, entry)) = matched_refresh {
runtime_state
.kv_set(
key.as_str(),
encode_kiro_prompt_cache_runtime_entry(entry),
Some(Duration::from_secs(entry.ttl_secs.max(1))),
)
.await?;
store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?;
}
let creation_tokens = last_breakpoint
@@ -176,21 +173,104 @@ async fn compute_kiro_prompt_cache_usage_with_runtime_state(
token_count: breakpoint.cumulative_tokens,
ttl_secs,
});
runtime_state
.kv_set(
key.as_str(),
encode_kiro_prompt_cache_runtime_entry(entry),
Some(Duration::from_secs(entry.ttl_secs.max(1))),
)
.await?;
store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?;
}
trim_kiro_prompt_cache_runtime_state(runtime_state, MAX_ENTRIES).await;
Ok(KiroPromptCacheUsage {
cache_creation_input_tokens: creation_tokens,
cache_read_input_tokens: matched_tokens,
})
}
async fn store_kiro_prompt_cache_runtime_entry(
runtime_state: &RuntimeState,
key: &str,
entry: KiroPromptCacheRuntimeEntry,
) -> Result<(), DataLayerError> {
let ttl = Duration::from_secs(entry.ttl_secs.max(1));
runtime_state
.kv_set(
key,
encode_kiro_prompt_cache_runtime_entry(entry),
Some(ttl),
)
.await?;
let expires_at_ms = current_unix_ms().saturating_add(entry.ttl_secs.saturating_mul(1000));
if let Err(err) = runtime_state
.score_set(KIRO_PROMPT_CACHE_INDEX_KEY, key, expires_at_ms as f64)
.await
{
warn!(
event_name = "kiro_simulated_cache_index_update_failed",
log_type = "event",
cache_key = %key,
error = ?err,
"failed to update Kiro simulated cache index; cache entry was persisted but cleanup may lag"
);
}
Ok(())
}
async fn trim_kiro_prompt_cache_runtime_state(runtime_state: &RuntimeState, max_entries: usize) {
if let Err(err) = runtime_state
.score_remove_by_score(KIRO_PROMPT_CACHE_INDEX_KEY, current_unix_ms() as f64)
.await
{
warn!(
event_name = "kiro_simulated_cache_index_expiry_trim_failed",
log_type = "event",
error = ?err,
"failed to trim expired Kiro simulated cache index entries"
);
return;
}
let Ok(index_len) = runtime_state.score_len(KIRO_PROMPT_CACHE_INDEX_KEY).await else {
return;
};
if index_len <= max_entries {
return;
}
let Ok(all_members) = runtime_state
.score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0)
.await
else {
return;
};
let trim_count = index_len.saturating_sub(max_entries);
if trim_count == 0 {
return;
}
let trimmed_members = all_members.into_iter().take(trim_count).collect::<Vec<_>>();
if let Err(err) = runtime_state.kv_delete_many(&trimmed_members).await {
warn!(
event_name = "kiro_simulated_cache_kv_trim_failed",
log_type = "event",
error = ?err,
trim_count,
"failed to delete trimmed Kiro simulated cache KV entries"
);
}
if let Err(err) = runtime_state
.score_remove_by_rank(KIRO_PROMPT_CACHE_INDEX_KEY, 0, trim_count as i64 - 1)
.await
{
warn!(
event_name = "kiro_simulated_cache_index_trim_failed",
log_type = "event",
error = ?err,
trim_count,
"failed to delete trimmed Kiro simulated cache index entries"
);
}
}
fn parse_kiro_prompt_cache_runtime_entry(value: &str) -> Option<KiroPromptCacheRuntimeEntry> {
serde_json::from_str::<KiroPromptCacheRuntimeEntry>(value)
.ok()
@@ -987,6 +1067,66 @@ mod tests {
assert!(second.cache_read_input_tokens > 0);
}
#[tokio::test]
async fn runtime_state_tracker_trims_oldest_entries_to_capacity() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
let now_ms = current_unix_ms();
let keys = [
"kiro:prompt-cache:test-oldest".to_string(),
"kiro:prompt-cache:test-middle".to_string(),
"kiro:prompt-cache:test-newest".to_string(),
];
for (index, key) in keys.iter().enumerate() {
runtime
.kv_set(
key,
encode_kiro_prompt_cache_runtime_entry(KiroPromptCacheRuntimeEntry {
token_count: 100 + index as u64,
ttl_secs: 120,
}),
Some(Duration::from_secs(120)),
)
.await
.expect("cache entry should store");
runtime
.score_set(
KIRO_PROMPT_CACHE_INDEX_KEY,
key,
now_ms.saturating_add(60_000 + index as u64 * 1_000) as f64,
)
.await
.expect("cache index should store");
}
trim_kiro_prompt_cache_runtime_state(&runtime, 2).await;
assert_eq!(
runtime
.kv_get(&keys[0])
.await
.expect("oldest entry should read"),
None
);
assert!(runtime
.kv_get(&keys[1])
.await
.expect("middle entry should read")
.is_some());
assert!(runtime
.kv_get(&keys[2])
.await
.expect("newest entry should read")
.is_some());
assert_eq!(
runtime
.score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0)
.await
.expect("cache index should read"),
vec![keys[1].clone(), keys[2].clone()]
);
}
#[test]
fn tracker_refreshes_cached_prefix_ttl_on_read() {
let base = serde_json::json!({
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+58 -25
View File
@@ -65,7 +65,7 @@ use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image
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,
compute_kiro_prompt_cache_usage, estimate_kiro_prompt_input_tokens, kiro_prompt_cache_tracker,
compute_kiro_prompt_cache_usage, estimate_kiro_prompt_input_tokens,
kiro_simulated_cache_enabled_from_provider_config,
kiro_simulated_cache_enabled_from_report_context, KiroPromptCacheUsage,
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD,
@@ -499,7 +499,8 @@ fn kiro_cache_usage_from_report_context(report_context: &Value) -> Option<KiroPr
.and_then(kiro_cache_usage_from_context_object)
}
fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
async fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
summary: &mut Option<ExecutionStreamTerminalSummary>,
@@ -577,8 +578,12 @@ fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
return;
};
let cache_usage = kiro_prompt_cache_tracker()
.compute_and_update(kiro_stream_cache_credential_id(plan), &profile);
let cache_usage = compute_kiro_prompt_cache_usage(
state.runtime_state(),
kiro_stream_cache_credential_id(plan),
&profile,
)
.await;
if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 {
return;
}
@@ -3705,10 +3710,12 @@ async fn execute_stream_from_frame_stream(
}
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
&mut stream_terminal_summary,
);
)
.await;
let requires_observed_terminal_event = stream_requires_observed_terminal_event(
plan_for_report.provider_api_format.as_str(),
stream_usage_report_context.as_ref(),
@@ -3951,6 +3958,10 @@ mod tests {
.with_execution_runtime_candidate(true)
}
fn test_state() -> AppState {
AppState::new().expect("gateway state should build")
}
#[test]
fn detects_client_visible_sse_terminal_events() {
assert!(stream_chunk_contains_sse_done(b"data: [DONE]\n\n"));
@@ -4129,8 +4140,8 @@ mod tests {
));
}
#[test]
fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() {
#[tokio::test]
async fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() {
let request_body = json!({
"model": "claude-opus-4-7",
"system": [
@@ -4178,6 +4189,7 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = test_state();
let mut first_summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4188,10 +4200,12 @@ mod tests {
..ExecutionStreamTerminalSummary::default()
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut first_summary,
);
)
.await;
let first_usage = first_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.as_ref())
@@ -4208,10 +4222,12 @@ mod tests {
..ExecutionStreamTerminalSummary::default()
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut second_summary,
);
)
.await;
let second_usage = second_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.as_ref())
@@ -4222,8 +4238,8 @@ mod tests {
assert_eq!(second_usage.output_tokens, 19);
}
#[test]
fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() {
#[tokio::test]
async fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() {
let first_request_body = json!({
"model": "claude-sonnet-4.6",
"messages": [{
@@ -4289,6 +4305,7 @@ mod tests {
"original_request_body": second_request_body,
"kiro_simulated_cache_enabled": true,
});
let state = test_state();
let mut first_summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4299,10 +4316,12 @@ mod tests {
..ExecutionStreamTerminalSummary::default()
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&first_report_context),
&mut first_summary,
);
)
.await;
let first_usage = first_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.as_ref())
@@ -4319,10 +4338,12 @@ mod tests {
..ExecutionStreamTerminalSummary::default()
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&second_report_context),
&mut second_summary,
);
)
.await;
let second_usage = second_summary
.as_ref()
.and_then(|summary| summary.standardized_usage.as_ref())
@@ -4335,8 +4356,8 @@ mod tests {
assert_eq!(second_usage.output_tokens, 19);
}
#[test]
fn kiro_stream_summary_seeds_input_tokens_without_cache_control() {
#[tokio::test]
async fn kiro_stream_summary_seeds_input_tokens_without_cache_control() {
let request_body = json!({
"model": "claude-opus-4-7",
"system": [
@@ -4382,6 +4403,7 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = test_state();
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4393,10 +4415,12 @@ mod tests {
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut summary,
);
)
.await;
let usage = summary
.as_ref()
@@ -4409,8 +4433,8 @@ mod tests {
assert_eq!(usage.output_tokens, 13);
}
#[test]
fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() {
#[tokio::test]
async fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() {
let request_body = json!({
"model": "claude-opus-4-7",
"system": [
@@ -4456,6 +4480,7 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = test_state();
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4468,10 +4493,12 @@ mod tests {
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut summary,
);
)
.await;
let usage = summary
.as_ref()
@@ -4484,8 +4511,8 @@ mod tests {
assert_eq!(usage.output_tokens, 23);
}
#[test]
fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() {
#[tokio::test]
async fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() {
let request_body = json!({
"model": "claude-opus-4-7",
"system": [
@@ -4532,6 +4559,7 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = test_state();
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4545,10 +4573,12 @@ mod tests {
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut summary,
);
)
.await;
let usage = summary
.as_ref()
@@ -4561,8 +4591,8 @@ mod tests {
assert_eq!(usage.output_tokens, 23);
}
#[test]
fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() {
#[tokio::test]
async fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() {
let request_body = json!({
"model": "claude-opus-4-7",
"messages": [
@@ -4610,6 +4640,7 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = test_state();
let mut summary = Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(StandardizedUsage {
@@ -4623,10 +4654,12 @@ mod tests {
});
maybe_apply_kiro_prompt_cache_usage_to_stream_summary(
&state,
&plan,
Some(&report_context),
&mut summary,
);
)
.await;
let usage = summary
.as_ref()