i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(kiro): simulate prompt cache usage accounting

Browse Source
This commit is contained in:
mayrain
2026-05-19 10:17:24 +08:00
parent c6a408d4e8
commit b4d17a392a
11 changed files with 2034 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
crates/aether-ai-formats/src/provider_compat/kiro_stream.rs
View File

@@ -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
}

1
crates/aether-ai-formats/src/provider_compat/kiro_stream/state.rs
View File

@@ -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>,

11
crates/aether-ai-formats/src/provider_compat/kiro_stream/stream/state/finalize.rs
View File

@@ -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)
}

18
crates/aether-ai-formats/src/provider_compat/kiro_stream/stream/state/lifecycle.rs
View File

@@ -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 {

78
crates/aether-ai-formats/src/provider_compat/kiro_stream/stream/tests.rs
View File

@@ -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);

27
crates/aether-data/src/repository/usage/postgres/queries/upsert_sql.sql
View File

@@ -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,

4
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -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",