Aether/crates/aether-ai/serving/src/dto.rs

use std::collections::BTreeMap;
use std::fmt;

use aether_ai_formats::api::ExecutionRuntimeAuthContext;
use aether_contracts::{
    redact_url_for_debug, ExecutionPlan, ExecutionTimeouts, ProxySnapshot, ResolvedTransportProfile,
};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionStrategy {
    GatewayAffinityForward,
    RawPublicProxy,
    LocalSameFormat,
    LocalCrossFormat,
}

impl ExecutionStrategy {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::GatewayAffinityForward => "gateway_affinity_forward",
            Self::RawPublicProxy => "raw_public_proxy",
            Self::LocalSameFormat => "local_same_format",
            Self::LocalCrossFormat => "local_cross_format",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ConversionMode {
    None,
    RequestOnly,
    ResponseOnly,
    Bidirectional,
}

impl ConversionMode {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::None => "none",
            Self::RequestOnly => "request_only",
            Self::ResponseOnly => "response_only",
            Self::Bidirectional => "bidirectional",
        }
    }
}

/// Request/response adaptation applied after candidate selection.
///
/// This is independent from format conversion: a same-format request may be
/// byte-transparent or may intentionally apply provider compatibility edits.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AdaptationMode {
    NativeTransparent,
    SameFormatCompat,
    CrossFormat,
}

impl AdaptationMode {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::NativeTransparent => "native_transparent",
            Self::SameFormatCompat => "same_format_compat",
            Self::CrossFormat => "cross_format",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
pub struct AiRequestGzipPolicy {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub enabled: Option<bool>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub min_bytes: Option<usize>,
}

#[derive(Deserialize, Serialize)]
pub struct AiExecutionPlanPayload {
    pub action: String,
    #[serde(default)]
    pub plan_kind: Option<String>,
    #[serde(default)]
    pub plan: Option<ExecutionPlan>,
    #[serde(default)]
    pub report_kind: Option<String>,
    #[serde(default)]
    pub report_context: Option<serde_json::Value>,
    #[serde(default)]
    pub auth_context: Option<ExecutionRuntimeAuthContext>,
}

impl fmt::Debug for AiExecutionPlanPayload {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("AiExecutionPlanPayload")
            .field("action", &self.action)
            .field("plan_kind", &self.plan_kind)
            .field("plan", &self.plan)
            .field("report_kind", &self.report_kind)
            .field("has_report_context", &self.report_context.is_some())
            .field(
                "report_context_bytes",
                &json_value_len(&self.report_context),
            )
            .field("has_auth_context", &self.auth_context.is_some())
            .finish()
    }
}

#[derive(Clone, Deserialize, Serialize)]
pub struct AiExecutionDecision {
    pub action: String,
    #[serde(default)]
    pub decision_kind: Option<String>,
    #[serde(default)]
    pub execution_strategy: Option<String>,
    #[serde(default)]
    pub conversion_mode: Option<String>,
    #[serde(default)]
    pub request_id: Option<String>,
    #[serde(default)]
    pub candidate_id: Option<String>,
    #[serde(default)]
    pub provider_name: Option<String>,
    #[serde(default)]
    pub provider_type: Option<String>,
    #[serde(default)]
    pub provider_id: Option<String>,
    #[serde(default)]
    pub endpoint_id: Option<String>,
    #[serde(default)]
    pub key_id: Option<String>,
    #[serde(default)]
    pub upstream_base_url: Option<String>,
    #[serde(default)]
    pub upstream_url: Option<String>,
    #[serde(default)]
    pub provider_request_method: Option<String>,
    #[serde(default)]
    pub auth_header: Option<String>,
    #[serde(default)]
    pub auth_value: Option<String>,
    #[serde(default)]
    pub provider_api_format: Option<String>,
    #[serde(default)]
    pub client_api_format: Option<String>,
    #[serde(default)]
    pub provider_contract: Option<String>,
    #[serde(default)]
    pub client_contract: Option<String>,
    #[serde(default)]
    pub model_name: Option<String>,
    #[serde(default)]
    pub mapped_model: Option<String>,
    #[serde(default)]
    pub prompt_cache_key: Option<String>,
    #[serde(default)]
    pub extra_headers: BTreeMap<String, String>,
    #[serde(default)]
    pub provider_request_headers: BTreeMap<String, String>,
    #[serde(default)]
    pub provider_request_body: Option<serde_json::Value>,
    #[serde(default)]
    pub provider_request_body_base64: Option<String>,
    #[serde(default)]
    pub content_type: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub content_encoding: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub request_gzip: Option<AiRequestGzipPolicy>,
    #[serde(default)]
    pub proxy: Option<ProxySnapshot>,
    #[serde(default)]
    pub transport_profile: Option<ResolvedTransportProfile>,
    #[serde(default)]
    pub timeouts: Option<ExecutionTimeouts>,
    #[serde(default)]
    pub upstream_is_stream: bool,
    #[serde(default)]
    pub report_kind: Option<String>,
    #[serde(default)]
    pub report_context: Option<serde_json::Value>,
    #[serde(default)]
    pub auth_context: Option<ExecutionRuntimeAuthContext>,
}

impl fmt::Debug for AiExecutionDecision {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = formatter.debug_struct("AiExecutionDecision");
        debug
            .field("action", &self.action)
            .field("decision_kind", &self.decision_kind)
            .field("execution_strategy", &self.execution_strategy)
            .field("conversion_mode", &self.conversion_mode)
            .field("request_id", &self.request_id)
            .field("candidate_id", &self.candidate_id)
            .field("provider_name", &self.provider_name)
            .field("provider_type", &self.provider_type)
            .field("provider_id", &self.provider_id)
            .field("endpoint_id", &self.endpoint_id)
            .field("key_id", &self.key_id)
            .field(
                "upstream_base_url",
                &self.upstream_base_url.as_deref().map(redact_url_for_debug),
            )
            .field(
                "upstream_url",
                &self.upstream_url.as_deref().map(redact_url_for_debug),
            )
            .field("provider_request_method", &self.provider_request_method)
            .field("auth_header", &self.auth_header)
            .field("has_auth_value", &self.auth_value.is_some())
            .field("auth_value_len", &self.auth_value.as_ref().map(String::len))
            .field("provider_api_format", &self.provider_api_format)
            .field("client_api_format", &self.client_api_format)
            .field("provider_contract", &self.provider_contract)
            .field("client_contract", &self.client_contract)
            .field("model_name", &self.model_name)
            .field("mapped_model", &self.mapped_model)
            .field("has_prompt_cache_key", &self.prompt_cache_key.is_some())
            .field(
                "prompt_cache_key_len",
                &self.prompt_cache_key.as_ref().map(String::len),
            )
            .field(
                "extra_header_names",
                &self.extra_headers.keys().collect::<Vec<_>>(),
            )
            .field(
                "provider_request_header_names",
                &self.provider_request_headers.keys().collect::<Vec<_>>(),
            )
            .field(
                "has_provider_request_body",
                &self.provider_request_body.is_some(),
            )
            .field(
                "provider_request_body_bytes",
                &json_value_len(&self.provider_request_body),
            )
            .field(
                "provider_request_body_base64_len",
                &self.provider_request_body_base64.as_ref().map(String::len),
            )
            .field("content_type", &self.content_type)
            .field("content_encoding", &self.content_encoding)
            .field("request_gzip", &self.request_gzip)
            .field("proxy", &self.proxy)
            .field("transport_profile", &self.transport_profile)
            .field("timeouts", &self.timeouts)
            .field("upstream_is_stream", &self.upstream_is_stream)
            .field("report_kind", &self.report_kind)
            .field("has_report_context", &self.report_context.is_some())
            .field(
                "report_context_bytes",
                &json_value_len(&self.report_context),
            )
            .field("has_auth_context", &self.auth_context.is_some())
            .finish()
    }
}

#[derive(Clone)]
pub struct AiSyncAttempt {
    pub plan: ExecutionPlan,
    pub report_kind: Option<String>,
    pub report_context: Option<serde_json::Value>,
}

impl fmt::Debug for AiSyncAttempt {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("AiSyncAttempt")
            .field("plan", &self.plan)
            .field("report_kind", &self.report_kind)
            .field("has_report_context", &self.report_context.is_some())
            .field(
                "report_context_bytes",
                &json_value_len(&self.report_context),
            )
            .finish()
    }
}

#[derive(Clone)]
pub struct AiStreamAttempt {
    pub plan: ExecutionPlan,
    pub report_kind: Option<String>,
    pub report_context: Option<serde_json::Value>,
}

impl fmt::Debug for AiStreamAttempt {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("AiStreamAttempt")
            .field("plan", &self.plan)
            .field("report_kind", &self.report_kind)
            .field("has_report_context", &self.report_context.is_some())
            .field(
                "report_context_bytes",
                &json_value_len(&self.report_context),
            )
            .finish()
    }
}

fn json_value_len(value: &Option<serde_json::Value>) -> Option<usize> {
    value
        .as_ref()
        .and_then(|value| serde_json::to_vec(value).ok().map(|bytes| bytes.len()))
}

pub fn augment_sync_report_context(
    report_context: Option<serde_json::Value>,
    provider_request_headers: &BTreeMap<String, String>,
    _provider_request_body: &serde_json::Value,
) -> serde_json::Result<Option<serde_json::Value>> {
    let mut report_context = match report_context {
        Some(serde_json::Value::Object(map)) => map,
        Some(_) => serde_json::Map::new(),
        None => serde_json::Map::new(),
    };

    report_context.insert(
        "provider_request_headers".to_string(),
        serde_json::to_value(provider_request_headers)?,
    );

    Ok(Some(serde_json::Value::Object(report_context)))
}

fn decision_has_exact_provider_request(payload: &AiExecutionDecision) -> bool {
    !payload.provider_request_headers.is_empty()
        && (payload.provider_request_body.is_some()
            || payload
                .provider_request_body_base64
                .as_ref()
                .map(|value| !value.trim().is_empty())
                .unwrap_or(false))
}

pub fn generic_decision_missing_exact_provider_request(payload: &AiExecutionDecision) -> bool {
    !decision_has_exact_provider_request(payload)
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::{
        augment_sync_report_context, generic_decision_missing_exact_provider_request,
        AiExecutionDecision,
    };

    #[test]
    fn generic_decision_detects_missing_exact_provider_request() {
        let payload = AiExecutionDecision {
            action: "local".to_string(),
            decision_kind: Some("sync".to_string()),
            execution_strategy: None,
            conversion_mode: None,
            request_id: None,
            candidate_id: None,
            provider_name: None,
            provider_type: None,
            provider_id: None,
            endpoint_id: None,
            key_id: None,
            upstream_base_url: None,
            upstream_url: None,
            provider_request_method: None,
            auth_header: None,
            auth_value: None,
            provider_api_format: None,
            client_api_format: None,
            provider_contract: None,
            client_contract: None,
            model_name: None,
            mapped_model: None,
            prompt_cache_key: None,
            extra_headers: Default::default(),
            provider_request_headers: Default::default(),
            provider_request_body: None,
            provider_request_body_base64: None,
            content_type: None,
            content_encoding: None,
            request_gzip: None,
            proxy: None,
            transport_profile: None,
            timeouts: None,
            upstream_is_stream: false,
            report_kind: None,
            report_context: None,
            auth_context: None,
        };

        assert!(generic_decision_missing_exact_provider_request(&payload));
    }

    #[test]
    fn augment_sync_report_context_attaches_provider_request_headers_only() {
        let report_context = augment_sync_report_context(
            Some(serde_json::json!({"trace_id": "abc"})),
            &BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
            &serde_json::json!({"model": "gpt-5"}),
        )
        .expect("context should serialize")
        .expect("context should exist");

        assert_eq!(
            report_context["provider_request_headers"]["content-type"],
            "application/json"
        );
        assert!(
            report_context.get("provider_request_body").is_none(),
            "provider request body should not be copied into report context"
        );
    }
}