Aether/crates/aether-contracts/src/plan.rs

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

use serde::{Deserialize, Serialize};
use serde_json::Value;

pub const EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER: &str = "x-aether-execution-follow-redirects";
pub const EXECUTION_REQUEST_HTTP1_ONLY_HEADER: &str = "x-aether-execution-http1-only";
pub const EXECUTION_RESPONSE_BODY_MODE_HEADER: &str = "x-aether-execution-response-body-mode";
pub const MAX_EXECUTION_REQUEST_TIMEOUT_SECS: u64 = 1_200;
pub const MAX_EXECUTION_REQUEST_TIMEOUT_MS: u64 = MAX_EXECUTION_REQUEST_TIMEOUT_SECS * 1_000;
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS: u64 = 300;
pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 =
    MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS * 1_000;

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionResponseBodyMode {
    #[default]
    StructuredJson,
    PreserveBytes,
}

impl ExecutionResponseBodyMode {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::StructuredJson => "structured_json",
            Self::PreserveBytes => "preserve_bytes",
        }
    }

    pub fn from_header_value(value: Option<&str>) -> Self {
        match value.map(str::trim) {
            Some(value) if value.eq_ignore_ascii_case(Self::PreserveBytes.as_str()) => {
                Self::PreserveBytes
            }
            _ => Self::StructuredJson,
        }
    }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct ExecutionTimeouts {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub connect_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub first_byte_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub write_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pool_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total_ms: Option<u64>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq)]
pub struct RequestBody {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub json_body: Option<Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub body_bytes_b64: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub body_ref: Option<String>,
}

impl fmt::Debug for RequestBody {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RequestBody")
            .field("has_json_body", &self.json_body.is_some())
            .field(
                "json_body_bytes",
                &self
                    .json_body
                    .as_ref()
                    .and_then(|body| serde_json::to_vec(body).ok().map(|bytes| bytes.len())),
            )
            .field(
                "body_bytes_b64_len",
                &self.body_bytes_b64.as_ref().map(String::len),
            )
            .field("body_ref_len", &self.body_ref.as_ref().map(String::len))
            .finish()
    }
}

impl RequestBody {
    pub fn from_json(json_body: Value) -> Self {
        Self {
            json_body: Some(json_body),
            body_bytes_b64: None,
            body_ref: None,
        }
    }
}

#[derive(Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct ProxySnapshot {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub enabled: Option<bool>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub mode: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub node_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub label: Option<String>,
    #[serde(default, alias = "proxy_url", skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub extra: Option<Value>,
}

impl fmt::Debug for ProxySnapshot {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ProxySnapshot")
            .field("enabled", &self.enabled)
            .field("mode", &self.mode)
            .field("node_id", &self.node_id)
            .field("label", &self.label)
            .field("url", &self.url.as_deref().map(redact_url_for_debug))
            .field("has_extra", &self.extra.is_some())
            .finish()
    }
}

/// Reserved internal metadata key used to fence proxy-node mutations against
/// a node incarnation recreated under the same stable id.
pub const PROXY_NODE_TUNNEL_GENERATION_EXTRA_KEY: &str = "proxy_node_tunnel_generation";

pub const TRANSPORT_BACKEND_REQWEST_RUSTLS: &str = "reqwest_rustls";
pub const TRANSPORT_BACKEND_HYPER_RUSTLS: &str = "hyper_rustls";
pub const TRANSPORT_BACKEND_BROWSER_WREQ: &str = "browser_wreq";
pub const TRANSPORT_HTTP_MODE_AUTO: &str = "auto";
pub const TRANSPORT_HTTP_MODE_HTTP1_ONLY: &str = "http1_only";
pub const TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE: &str = "h2c_prior_knowledge";
pub const TRANSPORT_POOL_SCOPE_KEY: &str = "key";

#[derive(Clone, Serialize, Deserialize, PartialEq)]
#[serde(default)]
pub struct ResolvedTransportProfile {
    pub profile_id: String,
    pub backend: String,
    pub http_mode: String,
    pub pool_scope: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub header_fingerprint: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub extra: Option<Value>,
}

impl Default for ResolvedTransportProfile {
    fn default() -> Self {
        Self {
            profile_id: String::new(),
            backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.to_string(),
            http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
            pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
            header_fingerprint: None,
            extra: None,
        }
    }
}

impl fmt::Debug for ResolvedTransportProfile {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ResolvedTransportProfile")
            .field("profile_id", &self.profile_id)
            .field("backend", &self.backend)
            .field("http_mode", &self.http_mode)
            .field("pool_scope", &self.pool_scope)
            .field("has_header_fingerprint", &self.header_fingerprint.is_some())
            .field("has_extra", &self.extra.is_some())
            .finish()
    }
}

#[derive(Clone, Serialize, Deserialize, PartialEq)]
pub struct ExecutionPlan {
    pub request_id: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub candidate_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub provider_name: Option<String>,
    pub provider_id: String,
    pub endpoint_id: String,
    pub key_id: String,
    pub method: String,
    #[serde(alias = "upstream_url")]
    pub url: String,
    #[serde(default)]
    pub headers: BTreeMap<String, String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub content_type: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub content_encoding: Option<String>,
    pub body: RequestBody,
    /// Whether the upstream API uses a streaming response protocol.
    #[serde(default)]
    pub stream: bool,
    pub client_api_format: String,
    pub provider_api_format: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model_name: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub proxy: Option<ProxySnapshot>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub transport_profile: Option<ResolvedTransportProfile>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub timeouts: Option<ExecutionTimeouts>,
}

impl fmt::Debug for ExecutionPlan {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ExecutionPlan")
            .field("request_id", &self.request_id)
            .field("candidate_id", &self.candidate_id)
            .field("provider_name", &self.provider_name)
            .field("provider_id", &self.provider_id)
            .field("endpoint_id", &self.endpoint_id)
            .field("key_id", &self.key_id)
            .field("method", &self.method)
            .field("url", &redact_url_for_debug(&self.url))
            .field("header_names", &self.headers.keys().collect::<Vec<_>>())
            .field("content_type", &self.content_type)
            .field("content_encoding", &self.content_encoding)
            .field("body", &self.body)
            .field("stream", &self.stream)
            .field("client_api_format", &self.client_api_format)
            .field("provider_api_format", &self.provider_api_format)
            .field("model_name", &self.model_name)
            .field("proxy", &self.proxy)
            .field("transport_profile", &self.transport_profile)
            .field("timeouts", &self.timeouts)
            .finish()
    }
}

/// Return a bounded URL representation suitable for diagnostics.
///
/// URL userinfo, query parameters, and fragments are never emitted because
/// providers commonly put API keys or OAuth tokens in those locations. An
/// unparsable URL is represented only by its length rather than echoing input.
pub fn redact_url_for_debug(raw: &str) -> String {
    const MAX_DEBUG_URL_CHARS: usize = 512;
    let raw = raw.trim();
    if raw.is_empty() {
        return String::new();
    }
    let Ok(mut url) = url::Url::parse(raw) else {
        return format!("[invalid-url len={}]", raw.len());
    };
    let _ = url.set_username("");
    let _ = url.set_password(None);
    url.set_query(None);
    url.set_fragment(None);
    let rendered = url.to_string();
    if rendered.chars().count() <= MAX_DEBUG_URL_CHARS {
        rendered
    } else {
        let prefix = rendered
            .chars()
            .take(MAX_DEBUG_URL_CHARS - 3)
            .collect::<String>();
        format!("{prefix}...")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn response_body_mode_is_fail_closed_to_structured_json() {
        assert_eq!(
            ExecutionResponseBodyMode::from_header_value(Some(" preserve_bytes ")),
            ExecutionResponseBodyMode::PreserveBytes
        );
        assert_eq!(
            ExecutionResponseBodyMode::from_header_value(Some("unexpected")),
            ExecutionResponseBodyMode::StructuredJson
        );
        assert_eq!(
            ExecutionResponseBodyMode::from_header_value(None),
            ExecutionResponseBodyMode::StructuredJson
        );
    }

    #[test]
    fn serializes_plan_with_json_body() {
        let plan = ExecutionPlan {
            request_id: "req_123".into(),
            candidate_id: Some("cand_123".into()),
            provider_name: Some("openai".into()),
            provider_id: "prov_123".into(),
            endpoint_id: "ep_123".into(),
            key_id: "key_123".into(),
            method: "POST".into(),
            url: "https://example.com/v1/chat/completions".into(),
            headers: BTreeMap::from([("authorization".into(), "Bearer test".into())]),
            content_type: Some("application/json".into()),
            content_encoding: Some("gzip".into()),
            body: RequestBody::from_json(serde_json::json!({"model":"gpt-test"})),
            stream: true,
            client_api_format: "openai:chat".into(),
            provider_api_format: "openai:chat".into(),
            model_name: Some("gpt-test".into()),
            proxy: None,
            transport_profile: None,
            timeouts: Some(ExecutionTimeouts {
                connect_ms: Some(30_000),
                read_ms: Some(3_600_000),
                first_byte_ms: Some(30_000),
                ..ExecutionTimeouts::default()
            }),
        };

        let raw = serde_json::to_value(&plan).expect("plan should serialize");
        assert_eq!(raw["body"]["json_body"]["model"], "gpt-test");
        assert_eq!(raw["content_encoding"], "gzip");
        assert_eq!(raw["stream"], true);
    }

    #[test]
    fn deserializes_python_control_plane_plan_shape() {
        let raw = serde_json::json!({
            "request_id": "req-1",
            "candidate_id": null,
            "provider_name": "openai",
            "provider_id": "prov-1",
            "endpoint_id": "ep-1",
            "key_id": "key-1",
            "method": "POST",
            "url": "https://example.com/v1/chat/completions",
            "headers": {"content-type": "application/json"},
            "content_encoding": "gzip",
            "body": {"json_body": {"model": "gpt-4.1"}},
            "stream": false,
            "provider_api_format": "openai:chat",
            "client_api_format": "openai:chat",
            "model_name": "gpt-4.1",
            "proxy": {
                "enabled": true,
                "mode": "direct",
                "label": "no-proxy",
                "url": "http://proxy.internal"
            },
            "timeouts": {
                "connect_ms": 10000,
                "read_ms": 30000,
                "write_ms": 30000,
                "pool_ms": 10000,
                "total_ms": 300000
            }
        });

        let plan: ExecutionPlan =
            serde_json::from_value(raw).expect("python payload should deserialize");
        assert_eq!(plan.url, "https://example.com/v1/chat/completions");
        assert_eq!(plan.candidate_id, None);
        assert_eq!(plan.provider_name.as_deref(), Some("openai"));
        assert_eq!(plan.model_name.as_deref(), Some("gpt-4.1"));
        assert_eq!(plan.content_encoding.as_deref(), Some("gzip"));
        assert_eq!(
            plan.proxy.as_ref().and_then(|proxy| proxy.url.as_deref()),
            Some("http://proxy.internal")
        );
        assert_eq!(
            plan.timeouts
                .as_ref()
                .and_then(|timeouts| timeouts.total_ms),
            Some(300_000)
        );
    }

    #[test]
    fn debug_redacts_plan_credentials_and_payloads() {
        let plan = ExecutionPlan {
            request_id: "req-debug".into(),
            candidate_id: Some("candidate-debug".into()),
            provider_name: Some("provider".into()),
            provider_id: "provider-id".into(),
            endpoint_id: "endpoint-id".into(),
            key_id: "key-id".into(),
            method: "POST".into(),
            url: "https://proxy-user:[email protected]/v1?api_key=url-secret#fragment-secret".into(),
            headers: BTreeMap::from([(
                "authorization".into(),
                "Bearer header-secret".into(),
            )]),
            content_type: Some("application/json".into()),
            content_encoding: None,
            body: RequestBody::from_json(serde_json::json!({
                "access_token": "body-secret",
                "prompt": "hello"
            })),
            stream: false,
            client_api_format: "openai:chat".into(),
            provider_api_format: "openai:chat".into(),
            model_name: Some("model".into()),
            proxy: Some(ProxySnapshot {
                enabled: Some(true),
                mode: Some("http".into()),
                node_id: Some("node".into()),
                label: None,
                url: Some("http://proxy-user:[email protected]?token=proxy-secret".into()),
                extra: Some(serde_json::json!({"credential": "extra-secret"})),
            }),
            transport_profile: Some(ResolvedTransportProfile {
                profile_id: "profile".into(),
                backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(),
                http_mode: TRANSPORT_HTTP_MODE_AUTO.into(),
                pool_scope: TRANSPORT_POOL_SCOPE_KEY.into(),
                header_fingerprint: Some(serde_json::json!({"authorization": "fingerprint-secret"})),
                extra: Some(serde_json::json!({"secret": "profile-secret"})),
            }),
            timeouts: None,
        };

        let debug = format!("{plan:?}");
        for secret in [
            "proxy-user",
            "proxy-password",
            "url-secret",
            "fragment-secret",
            "header-secret",
            "body-secret",
            "proxy-secret",
            "extra-secret",
            "fingerprint-secret",
            "profile-secret",
        ] {
            assert!(!debug.contains(secret), "debug leaked {secret}: {debug}");
        }
        assert!(debug.contains("header_names"));
        assert!(debug.contains("has_json_body"));
        assert!(debug.contains("profile"));
    }

    #[test]
    fn debug_url_redaction_fails_closed_for_invalid_urls() {
        let redacted = redact_url_for_debug("not a url?token=secret");
        assert!(!redacted.contains("secret"));
        assert!(redacted.starts_with("[invalid-url"));
    }
}