refactor ai serving modules and crates

crates/aether-ai-surfaces/src/finalize/stream_rewrite.rs

@@ -0,0 +1,356 @@
+use serde_json::Value;
+
+use crate::adaptation::kiro_stream::KiroToClaudeCliStreamState;
+use crate::adaptation::private_envelope::transform_provider_private_stream_line;
+use crate::adaptation::surfaces::{
+    provider_adaptation_should_unwrap_stream_envelope, KIRO_ENVELOPE_NAME,
+};
+use crate::finalize::openai_image_stream::OpenAiImageStreamState;
+use crate::finalize::standard::stream_core::StreamingStandardFormatMatrix;
+use crate::finalize::AiSurfaceFinalizeError;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum FinalizeStreamRewriteMode {
+    EnvelopeUnwrap,
+    OpenAiImage,
+    Standard,
+    KiroToClaudeCli,
+    KiroToClaudeCliThenStandard,
+}
+
+pub fn resolve_finalize_stream_rewrite_mode(
+    report_context: &Value,
+) -> Option<FinalizeStreamRewriteMode> {
+    let needs_conversion = report_context
+        .get("needs_conversion")
+        .and_then(Value::as_bool)
+        .unwrap_or(false);
+    let envelope_name = report_context
+        .get("envelope_name")
+        .and_then(Value::as_str)
+        .unwrap_or_default()
+        .trim()
+        .to_ascii_lowercase();
+    let provider_api_format = report_context
+        .get("provider_api_format")
+        .and_then(Value::as_str)
+        .unwrap_or_default()
+        .trim()
+        .to_ascii_lowercase();
+    let client_api_format = report_context
+        .get("client_api_format")
+        .and_then(Value::as_str)
+        .unwrap_or_default()
+        .trim()
+        .to_ascii_lowercase();
+
+    if needs_conversion
+        && envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
+        && provider_api_format == "claude:messages"
+    {
+        return supports_standard_stream_rewrite(
+            provider_api_format.as_str(),
+            client_api_format.as_str(),
+        )
+        .then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
+    }
+
+    if needs_conversion {
+        return supports_standard_stream_rewrite(
+            provider_api_format.as_str(),
+            client_api_format.as_str(),
+        )
+        .then_some(FinalizeStreamRewriteMode::Standard);
+    }
+
+    if provider_api_format == "openai:image" && client_api_format == "openai:image" {
+        return Some(FinalizeStreamRewriteMode::OpenAiImage);
+    }
+
+    if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
+        return (provider_api_format == "claude:messages"
+            && client_api_format == "claude:messages")
+            .then_some(FinalizeStreamRewriteMode::KiroToClaudeCli);
+    }
+
+    (provider_api_format == client_api_format
+        && provider_adaptation_should_unwrap_stream_envelope(
+            envelope_name.as_str(),
+            provider_api_format.as_str(),
+        ))
+    .then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
+}
+
+enum AiSurfaceStreamRewriteState {
+    EnvelopeUnwrap,
+    OpenAiImage(Box<OpenAiImageStreamState>),
+    Standard(Box<StreamingStandardFormatMatrix>),
+    KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
+    KiroToClaudeCliThenStandard {
+        kiro: Box<KiroToClaudeCliStreamState>,
+        standard: Box<StreamingStandardFormatMatrix>,
+    },
+}
+
+pub struct AiSurfaceStreamRewriter<'a> {
+    report_context: &'a Value,
+    buffered: Vec<u8>,
+    state: AiSurfaceStreamRewriteState,
+}
+
+pub fn maybe_build_ai_surface_stream_rewriter<'a>(
+    report_context: Option<&'a Value>,
+) -> Option<AiSurfaceStreamRewriter<'a>> {
+    let report_context = report_context?;
+    let state = match resolve_finalize_stream_rewrite_mode(report_context)? {
+        FinalizeStreamRewriteMode::EnvelopeUnwrap => AiSurfaceStreamRewriteState::EnvelopeUnwrap,
+        FinalizeStreamRewriteMode::OpenAiImage => {
+            AiSurfaceStreamRewriteState::OpenAiImage(Box::<OpenAiImageStreamState>::default())
+        }
+        FinalizeStreamRewriteMode::Standard => {
+            AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
+        }
+        FinalizeStreamRewriteMode::KiroToClaudeCli => AiSurfaceStreamRewriteState::KiroToClaudeCli(
+            Box::new(KiroToClaudeCliStreamState::new(report_context)),
+        ),
+        FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard => {
+            AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard {
+                kiro: Box::new(KiroToClaudeCliStreamState::new(report_context)),
+                standard: Box::<StreamingStandardFormatMatrix>::default(),
+            }
+        }
+    };
+
+    Some(AiSurfaceStreamRewriter {
+        report_context,
+        buffered: Vec::new(),
+        state,
+    })
+}
+
+impl AiSurfaceStreamRewriter<'_> {
+    pub fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
+        match &mut self.state {
+            AiSurfaceStreamRewriteState::OpenAiImage(state) => {
+                state.push_chunk(self.report_context, chunk)
+            }
+            AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
+                state.push_chunk(self.report_context, chunk)
+            }
+            AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
+                let claude_bytes = kiro.push_chunk(self.report_context, chunk)?;
+                transform_standard_bytes(standard, self.report_context, claude_bytes)
+            }
+            AiSurfaceStreamRewriteState::EnvelopeUnwrap
+            | AiSurfaceStreamRewriteState::Standard(_) => {
+                self.buffered.extend_from_slice(chunk);
+                let mut output = Vec::new();
+                while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
+                    let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
+                    output.extend(self.transform_line(line)?);
+                }
+                Ok(output)
+            }
+        }
+    }
+
+    pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
+        match &mut self.state {
+            AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context),
+            AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
+                state.finish(self.report_context)
+            }
+            AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
+                let mut output = transform_standard_bytes(
+                    standard,
+                    self.report_context,
+                    kiro.finish(self.report_context)?,
+                )?;
+                output.extend(standard.finish(self.report_context)?);
+                Ok(output)
+            }
+            AiSurfaceStreamRewriteState::EnvelopeUnwrap
+            | AiSurfaceStreamRewriteState::Standard(_) => {
+                if self.buffered.is_empty() {
+                    if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
+                        return state.finish(self.report_context);
+                    }
+                    return Ok(Vec::new());
+                }
+                let line = std::mem::take(&mut self.buffered);
+                let mut output = self.transform_line(line)?;
+                if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
+                    output.extend(state.finish(self.report_context)?);
+                }
+                Ok(output)
+            }
+        }
+    }
+
+    fn transform_line(&mut self, line: Vec<u8>) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
+        match &mut self.state {
+            AiSurfaceStreamRewriteState::EnvelopeUnwrap => {
+                transform_provider_private_stream_line(self.report_context, line)
+                    .map_err(AiSurfaceFinalizeError::from)
+            }
+            AiSurfaceStreamRewriteState::Standard(state) => {
+                transform_standard_line(state, self.report_context, line)
+            }
+            AiSurfaceStreamRewriteState::OpenAiImage(_)
+            | AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
+            | AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
+        }
+    }
+}
+
+fn transform_standard_bytes(
+    standard: &mut StreamingStandardFormatMatrix,
+    report_context: &Value,
+    bytes: Vec<u8>,
+) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
+    if bytes.is_empty() {
+        return Ok(Vec::new());
+    }
+    let mut output = Vec::new();
+    for line in bytes.split_inclusive(|byte| *byte == b'\n') {
+        output.extend(transform_standard_line(
+            standard,
+            report_context,
+            line.to_vec(),
+        )?);
+    }
+    Ok(output)
+}
+
+fn transform_standard_line(
+    standard: &mut StreamingStandardFormatMatrix,
+    report_context: &Value,
+    line: Vec<u8>,
+) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
+    let line = if should_unwrap_envelope(report_context) {
+        transform_provider_private_stream_line(report_context, line)?
+    } else {
+        line
+    };
+    if line.is_empty() {
+        return Ok(Vec::new());
+    }
+    standard.transform_line(report_context, line)
+}
+
+fn should_unwrap_envelope(report_context: &Value) -> bool {
+    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();
+    provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
+}
+
+fn supports_standard_stream_rewrite(provider_api_format: &str, client_api_format: &str) -> bool {
+    is_standard_provider_api_format(provider_api_format)
+        && (is_standard_chat_client_api_format(client_api_format)
+            || is_standard_cli_client_api_format(client_api_format))
+}
+
+fn is_standard_provider_api_format(api_format: &str) -> bool {
+    matches!(
+        aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
+        "openai:chat"
+            | "openai:responses"
+            | "openai:responses:compact"
+            | "claude:messages"
+            | "gemini:generate_content"
+    )
+}
+
+fn is_standard_chat_client_api_format(api_format: &str) -> bool {
+    matches!(
+        api_format,
+        "openai:chat" | "claude:messages" | "gemini:generate_content"
+    )
+}
+
+fn is_standard_cli_client_api_format(api_format: &str) -> bool {
+    matches!(
+        aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
+        "openai:responses"
+            | "openai:responses:compact"
+            | "claude:messages"
+            | "gemini:generate_content"
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
+
+    #[test]
+    fn resolves_standard_mode_for_cross_format_standard_streams() {
+        let report_context = json!({
+            "provider_api_format": "claude:messages",
+            "client_api_format": "openai:chat",
+            "needs_conversion": true,
+        });
+        assert_eq!(
+            resolve_finalize_stream_rewrite_mode(&report_context),
+            Some(FinalizeStreamRewriteMode::Standard)
+        );
+    }
+
+    #[test]
+    fn resolves_envelope_unwrap_for_same_format_private_envelopes() {
+        let report_context = json!({
+            "provider_api_format": "gemini:generate_content",
+            "client_api_format": "gemini:generate_content",
+            "envelope_name": "antigravity:v1internal",
+            "needs_conversion": false,
+        });
+        assert_eq!(
+            resolve_finalize_stream_rewrite_mode(&report_context),
+            Some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
+        );
+    }
+
+    #[test]
+    fn resolves_kiro_same_format_streams_to_kiro_mode() {
+        let report_context = json!({
+            "provider_api_format": "claude:messages",
+            "client_api_format": "claude:messages",
+            "envelope_name": "kiro:generateAssistantResponse",
+            "needs_conversion": false,
+        });
+        assert_eq!(
+            resolve_finalize_stream_rewrite_mode(&report_context),
+            Some(FinalizeStreamRewriteMode::KiroToClaudeCli)
+        );
+    }
+
+    #[test]
+    fn rejects_unsupported_non_conversion_streams() {
+        let report_context = json!({
+            "provider_api_format": "openai:chat",
+            "client_api_format": "openai:chat",
+            "needs_conversion": false,
+        });
+        assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
+    }
+
+    #[test]
+    fn resolves_openai_image_mode_for_same_format_image_streams() {
+        let report_context = json!({
+            "provider_api_format": "openai:image",
+            "client_api_format": "openai:image",
+            "needs_conversion": false,
+        });
+        assert_eq!(
+            resolve_finalize_stream_rewrite_mode(&report_context),
+            Some(FinalizeStreamRewriteMode::OpenAiImage)
+        );
+    }
+}