Fix Kiro cache usage reporting

This commit is contained in:
elky
2026-06-02 23:06:29 +08:00
parent 08ac9c5c58
commit 746af0d93e
3 changed files with 15 additions and 3 deletions
@@ -841,7 +841,6 @@ pub(crate) async fn execute_execution_runtime_stream(
) -> Result<Option<Response<Body>>, GatewayError> { ) -> Result<Option<Response<Body>>, GatewayError> {
let stream_started_at = Instant::now(); let stream_started_at = Instant::now();
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await; 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 lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let request_candidate_status_snapshot = let request_candidate_status_snapshot =
snapshot_local_request_candidate_status(&plan, report_context.as_ref()); snapshot_local_request_candidate_status(&plan, report_context.as_ref());
@@ -1991,9 +1990,11 @@ async fn execute_stream_from_frame_stream(
}; };
let mut report_context = let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers); attach_provider_response_headers_to_report_context(report_context, &headers);
seed_kiro_report_context_input_tokens(&plan, &mut report_context);
if status_code == 200 { if status_code == 200 {
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
seed_kiro_report_context_input_tokens(&plan, &mut report_context);
}
seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
} }
let mut buffered_frames = VecDeque::new(); let mut buffered_frames = VecDeque::new();
@@ -2117,8 +2117,10 @@ async fn execute_execution_runtime_sync_impl(
let mut report_context = let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers); attach_provider_response_headers_to_report_context(report_context, &headers);
if (200..300).contains(&status_code) { if (200..300).contains(&status_code) {
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await; seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
}
seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
} }
let mut client_headers = headers.clone(); let mut client_headers = headers.clone();
@@ -278,6 +278,9 @@ pub(crate) fn normalize_admin_user_api_formats(
if item.is_empty() { if item.is_empty() {
return Err("allowed_api_formats 不能为空".to_string()); return Err("allowed_api_formats 不能为空".to_string());
} }
if !looks_like_admin_api_format_signature(item) {
return Err(format!("allowed_api_formats 格式无效: {item}"));
}
let Some(normalized_item) = crate::api::ai::normalize_admin_endpoint_signature(item) else { let Some(normalized_item) = crate::api::ai::normalize_admin_endpoint_signature(item) else {
return Err(format!("allowed_api_formats 格式无效: {item}")); return Err(format!("allowed_api_formats 格式无效: {item}"));
}; };
@@ -289,6 +292,12 @@ pub(crate) fn normalize_admin_user_api_formats(
Ok(Some(normalized)) Ok(Some(normalized))
} }
fn looks_like_admin_api_format_signature(value: &str) -> bool {
value
.split_once(':')
.is_some_and(|(family, kind)| !family.trim().is_empty() && !kind.trim().is_empty())
}
pub(crate) fn normalize_admin_user_ip_rules( pub(crate) fn normalize_admin_user_ip_rules(
value: Option<Vec<String>>, value: Option<Vec<String>>,
) -> Result<Option<Vec<String>>, String> { ) -> Result<Option<Vec<String>>, String> {