Merge remote-tracking branch 'origin/pr/614'

2026-10-07 18:07:47 +08:00 · 2026-06-02 22:10:41 +08:00
parent bce3bf2b6e d1b64b6748
commit 08ac9c5c58
6 changed files with 1300 additions and 138 deletions
@@ -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,
-    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,
@@ -402,7 +402,8 @@ async fn seed_kiro_simulated_cache_enabled(
    }
 }

-fn seed_kiro_report_context_prompt_cache_usage(
+async fn seed_kiro_report_context_prompt_cache_usage(
+    state: &AppState,
    plan: &ExecutionPlan,
    report_context: &mut Option<Value>,
 ) {
@@ -450,8 +451,12 @@ fn seed_kiro_report_context_prompt_cache_usage(
        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;
    }
@@ -494,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>,
@@ -572,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;
    }
@@ -1984,7 +1994,7 @@ async fn execute_stream_from_frame_stream(
    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);
+        seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
    }
    let mut buffered_frames = VecDeque::new();
    let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None;
@@ -3705,10 +3715,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(),
@@ -3952,6 +3964,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"));
@@ -4136,8 +4152,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": [
@@ -4185,6 +4201,7 @@ mod tests {
            transport_profile: None,
            timeouts: None,
        };
+        let state = test_state();

        let mut first_summary = Some(ExecutionStreamTerminalSummary {
            standardized_usage: Some(StandardizedUsage {
@@ -4195,10 +4212,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())
@@ -4215,10 +4234,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())
@@ -4229,8 +4250,126 @@ 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_reads_cached_prefix_within_prompt_cache_lookback_window() {
+        let first_request_body = json!({
+            "model": "claude-sonnet-4.6",
+            "messages": [{
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": "shared first turn ".repeat(600),
+                    "cache_control": {"type": "ephemeral"}
+                }]
+            }]
+        });
+        let mut second_messages = vec![json!({
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": "shared first turn ".repeat(600)
+            }]
+        })];
+        for index in 0..12 {
+            second_messages.push(json!({
+                "role": if index % 2 == 0 { "assistant" } else { "user" },
+                "content": format!("intermediate stream turn {index}")
+            }));
+        }
+        second_messages.push(json!({
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": "new tail turn ".repeat(600),
+                "cache_control": {"type": "ephemeral"}
+            }]
+        }));
+        let second_request_body = json!({
+            "model": "claude-sonnet-4.6",
+            "messages": second_messages
+        });
+        let plan = ExecutionPlan {
+            request_id: "req-kiro-cache-stream-long-tail".into(),
+            candidate_id: Some("cand-kiro-cache-stream-long-tail".into()),
+            provider_name: Some("Kiro".into()),
+            provider_id: "provider-kiro-cache-stream-long-tail".into(),
+            endpoint_id: "endpoint-kiro-cache-stream-long-tail".into(),
+            key_id: "key-kiro-cache-stream-long-tail".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-sonnet-4.6".into()),
+            proxy: None,
+            transport_profile: None,
+            timeouts: None,
+        };
+        let first_report_context = json!({
+            "original_request_body": first_request_body,
+            "kiro_simulated_cache_enabled": true,
+        });
+        let second_report_context = json!({
+            "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 {
+                input_tokens: 4_000,
+                output_tokens: 17,
+                ..StandardizedUsage::new()
+            }),
+            ..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())
+            .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: 8_000,
+                output_tokens: 19,
+                ..StandardizedUsage::new()
+            }),
+            ..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())
+            .expect("second usage should exist");
+        assert!(
+            second_usage.cache_read_tokens > 0,
+            "stream summary should reuse the far earlier cached prefix"
+        );
+        assert!(second_usage.cache_creation_tokens > 0);
+        assert_eq!(second_usage.output_tokens, 19);
+    }
+
+    #[tokio::test]
+    async fn kiro_stream_summary_seeds_input_tokens_without_cache_control() {
        let request_body = json!({
            "model": "claude-opus-4-7",
            "system": [
@@ -4276,6 +4415,7 @@ mod tests {
            transport_profile: None,
            timeouts: None,
        };
+        let state = test_state();

        let mut summary = Some(ExecutionStreamTerminalSummary {
            standardized_usage: Some(StandardizedUsage {
@@ -4287,10 +4427,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()
@@ -4303,8 +4445,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": [
@@ -4350,6 +4492,7 @@ mod tests {
            transport_profile: None,
            timeouts: None,
        };
+        let state = test_state();

        let mut summary = Some(ExecutionStreamTerminalSummary {
            standardized_usage: Some(StandardizedUsage {
@@ -4362,10 +4505,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()
@@ -4378,8 +4523,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": [
@@ -4426,6 +4571,7 @@ mod tests {
            transport_profile: None,
            timeouts: None,
        };
+        let state = test_state();

        let mut summary = Some(ExecutionStreamTerminalSummary {
            standardized_usage: Some(StandardizedUsage {
@@ -4439,10 +4585,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()
@@ -4455,8 +4603,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": [
@@ -4504,6 +4652,7 @@ mod tests {
            transport_profile: None,
            timeouts: None,
        };
+        let state = test_state();

        let mut summary = Some(ExecutionStreamTerminalSummary {
            standardized_usage: Some(StandardizedUsage {
@@ -4517,10 +4666,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()
@@ -4640,9 +4791,11 @@ mod tests {
            "original_request_body": request_body,
            "kiro_simulated_cache_enabled": true,
        }));
+        let state = AppState::new().expect("gateway state should build");

        super::seed_kiro_report_context_input_tokens(&plan, &mut report_context);
-        super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context);
+        super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context)
+            .await;

        let context = report_context.as_ref().expect("context should exist");
        assert!(context
@@ -4709,9 +4862,11 @@ mod tests {
        let mut report_context = Some(json!({
            "original_request_body": request_body,
        }));
+        let state = AppState::new().expect("gateway state should build");

        super::seed_kiro_report_context_input_tokens(&plan, &mut report_context);
-        super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context);
+        super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context)
+            .await;

        let context = report_context.as_ref().expect("context should exist");
        assert!(context