修复：Kiro 模拟缓存接入共享运行时

2026-10-09 02:47:45 +08:00 · 2026-05-31 22:13:45 +08:00
parent 9734be31cf
commit c2bcfab7d4
6 changed files with 734 additions and 45 deletions
@@ -41,6 +41,12 @@ use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
 use crate::control::GatewayControlDecision;
 use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
 use crate::execution_runtime::grok::maybe_execute_grok_sync;
+use crate::execution_runtime::kiro_cache::{
+    build_kiro_prompt_cache_profile, 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,
+};
 use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
 #[cfg(test)]
 use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
@@ -355,6 +361,180 @@ fn build_sync_report_payload(
    }
 }

+fn seed_kiro_sync_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_sync_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);
+    }
+}
+
+async fn seed_kiro_sync_report_context_prompt_cache_usage(
+    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 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 = compute_kiro_prompt_cache_usage(
+        state.runtime_state(),
+        kiro_sync_cache_credential_id(plan),
+        &profile,
+    )
+    .await;
+    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_sync_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 invalid_gemini_provider_success_message(
    plan: &ExecutionPlan,
    report_context: Option<&Value>,
@@ -1934,8 +2114,13 @@ async fn execute_execution_runtime_sync_impl(
    }
    let status_code = result.status_code;
    let has_body_bytes = body_base64.is_some();
-    let report_context =
+    let mut report_context =
        attach_provider_response_headers_to_report_context(report_context, &headers);
+    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_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
+    }
    let mut client_headers = headers.clone();
    apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
        .await?;
@@ -2504,6 +2689,45 @@ mod tests {
        plan
    }

+    fn test_kiro_sync_plan() -> ExecutionPlan {
+        ExecutionPlan {
+            request_id: "req-kiro-sync-cache-1".to_string(),
+            candidate_id: Some("candidate-kiro-sync-cache-1".to_string()),
+            provider_name: Some("Kiro".to_string()),
+            provider_id: "provider-kiro-sync-1".to_string(),
+            endpoint_id: "endpoint-kiro-sync-1".to_string(),
+            key_id: "key-kiro-sync-1".to_string(),
+            method: "POST".to_string(),
+            url: "https://kiro.example/generateAssistantResponse".to_string(),
+            headers: BTreeMap::new(),
+            content_type: Some("application/json".to_string()),
+            content_encoding: None,
+            body: aether_contracts::RequestBody::from_json(json!({
+                "model": "claude-sonnet-4",
+                "messages": [{"role": "user", "content": "hello kiro"}],
+            })),
+            stream: false,
+            client_api_format: "claude:messages".to_string(),
+            provider_api_format: "claude:messages".to_string(),
+            model_name: Some("claude-sonnet-4".to_string()),
+            proxy: None,
+            transport_profile: None,
+            timeouts: None,
+        }
+    }
+
+    fn test_kiro_sync_cacheable_request_body() -> serde_json::Value {
+        json!({
+            "model": "claude-sonnet-4",
+            "system": [{
+                "type": "text",
+                "text": format!("sync cacheable prompt {}", "cacheable prompt chunk ".repeat(300)),
+                "cache_control": {"type": "ephemeral"}
+            }],
+            "messages": [{"role": "user", "content": "reuse this Kiro prompt"}]
+        })
+    }
+
    #[test]
    fn invalid_gemini_provider_success_uses_plan_format_when_context_is_missing() {
        let plan = test_gemini_chat_plan();
@@ -2659,6 +2883,68 @@ mod tests {
        );
    }

+    #[test]
+    fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body() {
+        let plan = test_kiro_sync_plan();
+        let mut report_context = Some(json!({
+            "original_request_body": test_kiro_sync_cacheable_request_body(),
+        }));
+
+        seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
+
+        assert!(report_context
+            .as_ref()
+            .and_then(|value| value.get("input_tokens"))
+            .and_then(Value::as_u64)
+            .is_some_and(|tokens| tokens > 0));
+    }
+
+    #[tokio::test]
+    async fn kiro_sync_report_context_applies_prompt_cache_usage_from_tracker() {
+        let state = AppState::new().expect("gateway state should build");
+        let plan = test_kiro_sync_plan();
+
+        let mut first_report_context = Some(json!({
+            "original_request_body": test_kiro_sync_cacheable_request_body(),
+            "kiro_simulated_cache_enabled": true,
+        }));
+        seed_kiro_sync_report_context_input_tokens(&plan, &mut first_report_context);
+        seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut first_report_context)
+            .await;
+        let first_creation = first_report_context
+            .as_ref()
+            .and_then(|value| value.get("cache_creation_input_tokens"))
+            .and_then(Value::as_u64)
+            .unwrap_or_default();
+        let first_read = first_report_context
+            .as_ref()
+            .and_then(|value| value.get("cache_read_input_tokens"))
+            .and_then(Value::as_u64)
+            .unwrap_or_default();
+        assert!(first_creation > 0);
+        assert_eq!(first_read, 0);
+
+        let mut second_report_context = Some(json!({
+            "original_request_body": test_kiro_sync_cacheable_request_body(),
+            "kiro_simulated_cache_enabled": true,
+        }));
+        seed_kiro_sync_report_context_input_tokens(&plan, &mut second_report_context);
+        seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut second_report_context)
+            .await;
+        let second_creation = second_report_context
+            .as_ref()
+            .and_then(|value| value.get("cache_creation_input_tokens"))
+            .and_then(Value::as_u64)
+            .unwrap_or_default();
+        let second_read = second_report_context
+            .as_ref()
+            .and_then(|value| value.get("cache_read_input_tokens"))
+            .and_then(Value::as_u64)
+            .unwrap_or_default();
+        assert_eq!(second_creation, 0);
+        assert!(second_read > 0);
+    }
+
    #[tokio::test]
    async fn json_whitespace_heartbeat_stream_prefixes_final_json() {
        let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);