Fix native Antigravity stream envelope handling

2026-10-05 00:47:48 +08:00 · 2026-05-25 09:39:55 +08:00
parent 28c3a5dbe4
commit b347f1816d
13 changed files with 770 additions and 57 deletions
@@ -47,6 +47,18 @@ pub fn resolve_finalize_stream_rewrite_mode(
        .trim()
        .to_ascii_lowercase();

+    if !needs_conversion
+        && client_consumes_same_private_stream_envelope(
+            report_context,
+            envelope_name.as_str(),
+            provider_api_format.as_str(),
+            client_api_format.as_str(),
+        )
+    {
+        return model_directive_display_model_from_report_context(report_context)
+            .map(|_| FinalizeStreamRewriteMode::ModelDirectiveDisplay);
+    }
+
    if needs_conversion
        && envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
        && provider_api_format == "claude:messages"
@@ -107,6 +119,26 @@ pub fn resolve_finalize_stream_rewrite_mode(
    .then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
 }

+fn client_consumes_same_private_stream_envelope(
+    report_context: &Value,
+    envelope_name: &str,
+    provider_api_format: &str,
+    client_api_format: &str,
+) -> bool {
+    if envelope_name.is_empty()
+        || provider_api_format != client_api_format
+        || !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
+    {
+        return false;
+    }
+    report_context
+        .get("client_envelope_name")
+        .and_then(Value::as_str)
+        .is_some_and(|client_envelope_name| {
+            client_envelope_name.eq_ignore_ascii_case(envelope_name)
+        })
+}
+
 enum AiSurfaceStreamRewriteState {
    EnvelopeUnwrap,
    ModelDirectiveDisplay,
@@ -451,6 +483,45 @@ mod tests {
        );
    }

+    #[test]
+    fn resolves_no_rewriter_when_client_consumes_same_private_envelope() {
+        let report_context = json!({
+            "provider_api_format": "gemini:generate_content",
+            "client_api_format": "gemini:generate_content",
+            "envelope_name": "antigravity:v1internal",
+            "client_envelope_name": "antigravity:v1internal",
+            "needs_conversion": false,
+        });
+        assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
+        assert!(maybe_build_ai_surface_stream_rewriter(Some(&report_context)).is_none());
+    }
+
+    #[test]
+    fn native_private_envelope_client_keeps_response_wrapper_for_model_display_rewrite() {
+        let report_context = json!({
+            "provider_api_format": "gemini:generate_content",
+            "client_api_format": "gemini:generate_content",
+            "envelope_name": "antigravity:v1internal",
+            "client_envelope_name": "antigravity:v1internal",
+            "model": "gemini-2.5-pro-high",
+            "mapped_model": "gemini-2.5-pro",
+            "needs_conversion": false,
+        });
+        let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
+            .expect("display-model rewriter should exist");
+        let output = rewriter
+            .push_chunk(
+                b"data: {\"response\":{\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[]},\"responseId\":\"resp_native_123\"}\n\n",
+            )
+            .expect("rewrite should succeed");
+        let output = String::from_utf8(output).expect("output should be utf8");
+
+        assert!(output.contains("\"response\":"));
+        assert!(output.contains("\"responseId\":\"resp_native_123\""));
+        assert!(output.contains("\"modelVersion\":\"gemini-2.5-pro-high\""));
+        assert!(!output.contains("_v1internal_response_id"));
+    }
+
    #[test]
    fn resolves_kiro_same_format_streams_to_kiro_mode() {
        let report_context = json!({
@@ -47,6 +47,13 @@ pub fn normalize_provider_private_report_context(report_context: Option<&Value>)
        .get("provider_api_format")
        .and_then(Value::as_str)
        .unwrap_or_default();
+    if report_context_preserves_private_client_envelope(
+        report_context,
+        envelope_name,
+        provider_api_format,
+    ) {
+        return Some(report_context.clone());
+    }
    if provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none() {
        return Some(report_context.clone());
    }
@@ -64,6 +71,22 @@ pub fn normalize_provider_private_response_value(
    {
        return Some(data);
    }
+    let envelope_name = report_context
+        .get("envelope_name")
+        .and_then(Value::as_str)
+        .unwrap_or_default();
+    let provider_api_format = report_context
+        .get("provider_api_format")
+        .and_then(Value::as_str)
+        .unwrap_or_default();
+    if report_context_preserves_private_client_envelope(
+        report_context,
+        envelope_name,
+        provider_api_format,
+    ) {
+        return Some(data);
+    }
+
    let mut unwrapped = match report_context.get("envelope_name").and_then(Value::as_str) {
        Some(KIRO_ENVELOPE_NAME) => data,
        Some(GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME) => {
@@ -155,6 +178,13 @@ fn transform_provider_private_stream_line_with_event_state(
    if !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format) {
        return Ok(line);
    }
+    if report_context_preserves_private_client_envelope(
+        report_context,
+        envelope_name,
+        provider_api_format,
+    ) {
+        return Ok(line);
+    }
    if envelope_name == WINDSURF_ENVELOPE_NAME && looks_like_windsurf_error(&body) {
        return Ok(line);
    }
@@ -305,6 +335,13 @@ pub fn maybe_build_provider_private_stream_normalizer<'a>(
        .unwrap_or_default();
    let descriptor =
        provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format)?;
+    if report_context_preserves_private_client_envelope(
+        report_context,
+        envelope_name,
+        provider_api_format,
+    ) {
+        return None;
+    }
    let mode = if descriptor
        .envelope_name
        .eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
@@ -679,6 +716,24 @@ fn clear_private_envelope_context(report_context: &Value) -> Value {
    normalized
 }

+fn report_context_preserves_private_client_envelope(
+    report_context: &Value,
+    envelope_name: &str,
+    provider_api_format: &str,
+) -> bool {
+    if envelope_name.is_empty()
+        || provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none()
+    {
+        return false;
+    }
+    report_context
+        .get("client_envelope_name")
+        .and_then(Value::as_str)
+        .is_some_and(|client_envelope_name| {
+            client_envelope_name.eq_ignore_ascii_case(envelope_name)
+        })
+}
+
 fn local_finalize_response_model(report_context: &Value) -> &str {
    report_context
        .get("mapped_model")
@@ -997,6 +1052,20 @@ mod tests {
        assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
    }

+    #[test]
+    fn private_stream_normalizer_preserves_antigravity_native_client_envelope() {
+        let report_context = json!({
+            "has_envelope": true,
+            "provider_api_format": "gemini:generate_content",
+            "client_api_format": "gemini:generate_content",
+            "envelope_name": "antigravity:v1internal",
+            "client_envelope_name": "antigravity:v1internal",
+            "mapped_model": "claude-sonnet-4-5",
+        });
+
+        assert!(maybe_build_provider_private_stream_normalizer(Some(&report_context)).is_none());
+    }
+
    #[test]
    fn detects_sse_error_events_without_explicit_type_field() {
        let body = br#"event: error