refactor ai serving modules and crates

apps/aether-gateway/src/ai_serving/finalize/common.rs

@@ -0,0 +1,139 @@
+use std::collections::BTreeMap;
+
+use axum::body::Body;
+use axum::http::Response;
+use serde_json::Value;
+
+pub(crate) use crate::ai_serving::api::{
+    normalize_provider_private_response_value as unwrap_local_finalize_response_value,
+    provider_private_response_allows_sync_finalize as local_finalize_allows_envelope,
+};
+use crate::ai_serving::{
+    build_generated_tool_call_id,
+    build_local_success_background_report as build_local_success_background_report_impl,
+    build_local_success_conversion_background_report as build_local_success_conversion_background_report_impl,
+    canonicalize_tool_arguments,
+    prepare_local_success_response_parts as prepare_local_success_response_parts_impl,
+    GatewayControlDecision, LocalSyncReportParts,
+};
+use crate::api::response::build_client_response_from_parts;
+use crate::{usage::GatewaySyncReportRequest, GatewayError};
+
+pub(crate) struct LocalCoreSyncFinalizeOutcome {
+    pub(crate) response: Response<Body>,
+    pub(crate) background_report: Option<GatewaySyncReportRequest>,
+}
+
+fn build_local_success_response(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    status_code: u16,
+    body_bytes: Vec<u8>,
+    headers: BTreeMap<String, String>,
+) -> Result<Response<Body>, GatewayError> {
+    build_client_response_from_parts(
+        status_code,
+        &headers,
+        Body::from(body_bytes),
+        trace_id,
+        Some(decision),
+    )
+}
+
+fn surface_report_parts_from_gateway(payload: &GatewaySyncReportRequest) -> LocalSyncReportParts {
+    LocalSyncReportParts {
+        trace_id: payload.trace_id.clone(),
+        report_kind: payload.report_kind.clone(),
+        report_context: payload.report_context.clone(),
+        status_code: payload.status_code,
+        headers: payload.headers.clone(),
+        body_json: payload.body_json.clone(),
+        client_body_json: payload.client_body_json.clone(),
+        body_base64: payload.body_base64.clone(),
+    }
+}
+
+fn gateway_report_from_surface(
+    source: &GatewaySyncReportRequest,
+    report: LocalSyncReportParts,
+) -> GatewaySyncReportRequest {
+    GatewaySyncReportRequest {
+        trace_id: report.trace_id,
+        report_kind: report.report_kind,
+        report_context: report.report_context,
+        status_code: report.status_code,
+        headers: report.headers,
+        body_json: report.body_json,
+        client_body_json: report.client_body_json,
+        body_base64: report.body_base64,
+        telemetry: source.telemetry.clone(),
+    }
+}
+
+pub(crate) fn build_local_success_outcome(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+    body_json: Value,
+) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
+    let report_headers = payload.headers.clone();
+    let (body_bytes, response_headers) =
+        prepare_local_success_response_parts_impl(&payload.headers, &body_json)
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+    let surface_payload = surface_report_parts_from_gateway(payload);
+    let background_report =
+        build_local_success_background_report_impl(&surface_payload, body_json, report_headers)
+            .map(|report| gateway_report_from_surface(payload, report));
+    build_local_success_outcome_with_report(
+        trace_id,
+        decision,
+        payload.status_code,
+        body_bytes,
+        response_headers,
+        background_report,
+    )
+}
+
+pub(crate) fn build_local_success_outcome_with_report(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    status_code: u16,
+    body_bytes: Vec<u8>,
+    headers: BTreeMap<String, String>,
+    background_report: Option<GatewaySyncReportRequest>,
+) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
+    let response =
+        build_local_success_response(trace_id, decision, status_code, body_bytes, headers)?;
+    Ok(LocalCoreSyncFinalizeOutcome {
+        response,
+        background_report,
+    })
+}
+
+pub(crate) fn build_local_success_outcome_with_conversion_report(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+    client_body_json: Value,
+    provider_body_json: Value,
+) -> Result<LocalCoreSyncFinalizeOutcome, GatewayError> {
+    let (body_bytes, response_headers) =
+        prepare_local_success_response_parts_impl(&payload.headers, &client_body_json)
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+    let surface_payload = surface_report_parts_from_gateway(payload);
+    let report_payload = build_local_success_conversion_background_report_impl(
+        &surface_payload,
+        client_body_json,
+        provider_body_json,
+    )
+    .map(|report| gateway_report_from_surface(payload, report));
+
+    build_local_success_outcome_with_report(
+        trace_id,
+        decision,
+        payload.status_code,
+        body_bytes,
+        response_headers,
+        report_payload,
+    )
+}

apps/aether-gateway/src/ai_serving/finalize/internal/mod.rs

@@ -0,0 +1,44 @@
+use axum::body::Body;
+use axum::http::Response;
+use serde_json::Value;
+
+use crate::ai_serving::GatewayControlDecision;
+use crate::{usage::GatewaySyncReportRequest, GatewayError};
+
+#[path = "stream_rewrite.rs"]
+pub(crate) mod stream;
+#[path = "sync_finalize.rs"]
+pub(crate) mod sync;
+#[path = "sync_to_stream.rs"]
+pub(crate) mod sync_to_stream;
+
+pub(crate) use stream::LocalStreamRewriter;
+pub(crate) use sync::LocalCoreSyncFinalizeOutcome;
+pub(crate) use sync_to_stream::{
+    maybe_bridge_standard_sync_json_to_stream, SyncToStreamBridgeOutcome,
+};
+
+pub(crate) fn maybe_build_sync_finalize_outcome(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
+    sync::maybe_build_local_core_sync_finalize_response(trace_id, decision, payload)
+}
+
+pub(crate) fn maybe_compile_sync_finalize_response(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+) -> Result<Option<Response<Body>>, GatewayError> {
+    Ok(
+        maybe_build_sync_finalize_outcome(trace_id, decision, payload)?
+            .map(|outcome| outcome.response),
+    )
+}
+
+pub(crate) fn maybe_build_stream_response_rewriter(
+    report_context: Option<&Value>,
+) -> Option<LocalStreamRewriter<'_>> {
+    stream::maybe_build_local_stream_rewriter(report_context)
+}

apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs

@@ -0,0 +1,35 @@
+use serde_json::Value;
+
+use crate::ai_serving::{
+    maybe_build_ai_surface_stream_rewriter, AiSurfaceFinalizeError, AiSurfaceStreamRewriter,
+};
+use crate::GatewayError;
+
+pub(crate) struct LocalStreamRewriter<'a> {
+    inner: AiSurfaceStreamRewriter<'a>,
+}
+
+pub(crate) fn maybe_build_local_stream_rewriter<'a>(
+    report_context: Option<&'a Value>,
+) -> Option<LocalStreamRewriter<'a>> {
+    maybe_build_ai_surface_stream_rewriter(report_context)
+        .map(|inner| LocalStreamRewriter { inner })
+}
+
+impl LocalStreamRewriter<'_> {
+    pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
+        self.inner.push_chunk(chunk).map_err(map_surface_error)
+    }
+
+    pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
+        self.inner.finish().map_err(map_surface_error)
+    }
+}
+
+fn map_surface_error(error: AiSurfaceFinalizeError) -> GatewayError {
+    error.into()
+}
+
+#[cfg(test)]
+#[path = "../tests_stream.rs"]
+mod tests;

apps/aether-gateway/src/ai_serving/finalize/internal/sync_finalize.rs

@@ -0,0 +1,113 @@
+use crate::ai_serving::GatewayControlDecision;
+use crate::ai_serving::{build_generated_tool_call_id, canonicalize_tool_arguments};
+use crate::{usage::GatewaySyncReportRequest, GatewayError};
+
+pub(crate) use crate::ai_serving::finalize::common::{
+    build_local_success_outcome, build_local_success_outcome_with_conversion_report,
+    local_finalize_allows_envelope, unwrap_local_finalize_response_value,
+    LocalCoreSyncFinalizeOutcome,
+};
+pub(crate) use crate::ai_serving::finalize::standard::{
+    maybe_build_standard_sync_finalize_product_from_normalized_payload,
+    StandardSyncFinalizeNormalizedProduct,
+};
+pub(crate) use crate::ai_serving::{
+    aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
+    aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
+    maybe_build_openai_image_sync_finalize_product,
+};
+pub(crate) use crate::ai_serving::{
+    convert_claude_chat_response_to_openai_chat, convert_claude_response_to_openai_responses,
+    convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
+};
+
+pub(crate) fn maybe_build_local_core_sync_finalize_response(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
+    if let Some(outcome) =
+        maybe_build_local_openai_image_sync_finalize_response(trace_id, decision, payload)?
+    {
+        return Ok(Some(outcome));
+    }
+
+    let Some(normalized_payload) =
+        crate::ai_serving::adaptation::private_envelope::maybe_normalize_provider_private_sync_report_payload(payload)?
+    else {
+        return Ok(None);
+    };
+    let payload = &normalized_payload;
+    let Some(report_context) = payload.report_context.as_ref() else {
+        return Ok(None);
+    };
+    if !local_finalize_allows_envelope(report_context) {
+        return Ok(None);
+    }
+    let Some(product) = maybe_build_standard_sync_finalize_product_from_normalized_payload(
+        payload.report_kind.as_str(),
+        payload.status_code,
+        Some(report_context),
+        payload.body_json.as_ref(),
+        payload.body_base64.as_deref(),
+    )
+    .map_err(GatewayError::from)?
+    else {
+        return Ok(None);
+    };
+
+    match product {
+        StandardSyncFinalizeNormalizedProduct::SuccessBody(body_json) => {
+            let Some(body_json) = unwrap_local_finalize_response_value(body_json, report_context)
+            else {
+                return Ok(None);
+            };
+            Ok(Some(build_local_success_outcome(
+                trace_id, decision, payload, body_json,
+            )?))
+        }
+        StandardSyncFinalizeNormalizedProduct::CrossFormat(product) => {
+            let Some(provider_body_json) =
+                unwrap_local_finalize_response_value(product.provider_body_json, report_context)
+            else {
+                return Ok(None);
+            };
+            Ok(Some(build_local_success_outcome_with_conversion_report(
+                trace_id,
+                decision,
+                payload,
+                product.client_body_json,
+                provider_body_json,
+            )?))
+        }
+    }
+}
+
+fn maybe_build_local_openai_image_sync_finalize_response(
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    payload: &GatewaySyncReportRequest,
+) -> Result<Option<LocalCoreSyncFinalizeOutcome>, GatewayError> {
+    let Some(product) = maybe_build_openai_image_sync_finalize_product(
+        payload.report_kind.as_str(),
+        payload.status_code,
+        payload.report_context.as_ref(),
+        payload.body_base64.as_deref(),
+    )
+    .map_err(GatewayError::from)?
+    else {
+        return Ok(None);
+    };
+
+    Ok(Some(build_local_success_outcome_with_conversion_report(
+        trace_id,
+        decision,
+        payload,
+        product.client_body_json,
+        product.provider_body_json,
+    )?))
+}
+
+#[cfg(test)]
+#[path = "../tests_sync.rs"]
+mod tests;

apps/aether-gateway/src/ai_serving/finalize/internal/sync_to_stream.rs

@@ -0,0 +1,20 @@
+use serde_json::Value;
+
+use crate::GatewayError;
+
+pub(crate) use crate::ai_serving::pure::SyncToStreamBridgeOutcome;
+
+pub(crate) fn maybe_bridge_standard_sync_json_to_stream(
+    provider_body_json: &Value,
+    provider_api_format: &str,
+    client_api_format: &str,
+    report_context: Option<&Value>,
+) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
+    crate::ai_serving::pure::maybe_bridge_standard_sync_json_to_stream(
+        provider_body_json,
+        provider_api_format,
+        client_api_format,
+        report_context,
+    )
+    .map_err(GatewayError::from)
+}

apps/aether-gateway/src/ai_serving/finalize/mod.rs

@@ -0,0 +1,4 @@
+pub(crate) mod common;
+pub(crate) mod internal;
+pub(crate) mod sse;
+pub(crate) mod standard;

apps/aether-gateway/src/ai_serving/finalize/sse.rs

@@ -0,0 +1,15 @@
+use serde_json::Value;
+
+use crate::ai_serving::{
+    encode_done_sse, encode_json_sse as encode_json_sse_impl, map_claude_stop_reason,
+    AiSurfaceFinalizeError,
+};
+use crate::GatewayError;
+
+fn map_error(err: AiSurfaceFinalizeError) -> GatewayError {
+    err.into()
+}
+
+pub(crate) fn encode_json_sse(event: Option<&str>, value: &Value) -> Result<Vec<u8>, GatewayError> {
+    encode_json_sse_impl(event, value).map_err(map_error)
+}

apps/aether-gateway/src/ai_serving/finalize/standard/mod.rs

@@ -0,0 +1,19 @@
+//! Standard finalize surface for standard contract sync/stream compilation.
+
+pub(crate) use crate::ai_serving::{
+    aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
+    build_openai_responses_response, convert_claude_chat_response_to_openai_chat,
+    convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
+    convert_gemini_response_to_openai_responses, convert_openai_chat_response_to_claude_chat,
+    convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_responses,
+    convert_openai_responses_response_to_openai_chat, convert_standard_chat_response,
+    convert_standard_cli_response,
+    maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload,
+    maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload,
+    maybe_build_openai_responses_same_family_sync_body_from_normalized_payload,
+    maybe_build_standard_cross_format_sync_product,
+    maybe_build_standard_cross_format_sync_product_from_normalized_payload,
+    maybe_build_standard_same_format_sync_body_from_normalized_payload,
+    maybe_build_standard_sync_finalize_product_from_normalized_payload,
+    StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
+};