Aether/crates/aether-provider/transport/src/openai_image/mod.rs

use std::collections::BTreeMap;

use serde_json::Value;

use crate::auth::{build_passthrough_headers_with_auth, resolve_local_openai_bearer_auth};
use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
use crate::policy::local_standard_transport_unsupported_reason_with_network;
use crate::rules::apply_local_header_rules_with_request_headers;
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::build_openai_image_url;

#[derive(Debug, Clone, Copy)]
pub struct ProviderOpenAiImageHeadersInput<'a> {
    pub transport: &'a GatewayProviderTransportSnapshot,
    pub headers: &'a http::HeaderMap,
    pub auth_header: &'a str,
    pub auth_value: &'a str,
    pub accept: Option<&'a str>,
    pub header_rules: Option<&'a Value>,
    pub provider_request_body: &'a Value,
    pub original_request_body: &'a Value,
}

pub fn openai_image_transport_unsupported_reason(
    transport: &GatewayProviderTransportSnapshot,
    api_format: &str,
) -> Option<&'static str> {
    let reason = local_standard_transport_unsupported_reason_with_network(transport, api_format);
    if is_dedicated_openai_image_provider(transport)
        && matches!(
            reason,
            Some("transport_provider_type_unsupported")
                | Some("transport_oauth_resolution_unsupported")
        )
    {
        return None;
    }
    reason
}

fn is_dedicated_openai_image_provider(transport: &GatewayProviderTransportSnapshot) -> bool {
    transport
        .provider
        .provider_type
        .trim()
        .eq_ignore_ascii_case("chatgpt_web")
        || transport
            .provider
            .provider_type
            .trim()
            .eq_ignore_ascii_case("codex")
        || is_grok_provider_transport(transport)
}

pub fn resolve_openai_image_auth(
    transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
    if is_grok_provider_transport(transport) {
        return resolve_grok_session_auth(transport);
    }
    resolve_local_openai_bearer_auth(transport)
}

pub fn build_openai_image_upstream_url(
    transport: &GatewayProviderTransportSnapshot,
    request_path: Option<&str>,
    request_query: Option<&str>,
) -> String {
    build_openai_image_url(&transport.endpoint.base_url, request_path, request_query)
}

pub fn build_openai_image_headers(
    input: ProviderOpenAiImageHeadersInput<'_>,
) -> Option<BTreeMap<String, String>> {
    let mut provider_request_headers = build_passthrough_headers_with_auth(
        input.headers,
        input.auth_header,
        input.auth_value,
        &BTreeMap::new(),
    );
    provider_request_headers.insert("content-type".to_string(), "application/json".to_string());
    if let Some(accept) = input.accept {
        provider_request_headers.insert("accept".to_string(), accept.to_string());
    } else {
        provider_request_headers.remove("accept");
    }
    crate::apply_local_auth_config_header_overrides(
        &mut provider_request_headers,
        input.transport.key.decrypted_auth_config.as_deref(),
    );
    if !apply_local_header_rules_with_request_headers(
        &mut provider_request_headers,
        input.header_rules,
        &[input.auth_header, "content-type", "accept"],
        input.provider_request_body,
        Some(input.original_request_body),
        Some(input.headers),
    ) {
        return None;
    }
    Some(provider_request_headers)
}

#[cfg(test)]
mod tests {
    use http::HeaderMap;
    use serde_json::json;

    use super::{
        build_openai_image_headers, build_openai_image_upstream_url,
        openai_image_transport_unsupported_reason, ProviderOpenAiImageHeadersInput,
    };
    use crate::snapshot::{
        GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
        GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
    };

    fn sample_transport() -> GatewayProviderTransportSnapshot {
        GatewayProviderTransportSnapshot {
            provider: GatewayProviderTransportProvider {
                id: "provider-1".to_string(),
                name: "Provider".to_string(),
                provider_type: "codex".to_string(),
                website: None,
                is_active: true,
                keep_priority_on_conversion: false,
                enable_format_conversion: false,
                concurrent_limit: None,
                max_retries: None,
                proxy: None,
                request_timeout_secs: None,
                stream_first_byte_timeout_secs: None,
                config: None,
            },
            endpoint: GatewayProviderTransportEndpoint {
                id: "endpoint-1".to_string(),
                provider_id: "provider-1".to_string(),
                api_format: "openai:image".to_string(),
                api_family: None,
                endpoint_kind: None,
                is_active: true,
                base_url: "https://api.openai.com/v1".to_string(),
                header_rules: None,
                body_rules: None,
                max_retries: None,
                custom_path: None,
                config: None,
                format_acceptance_config: None,
                proxy: None,
            },
            key: GatewayProviderTransportKey {
                id: "key-1".to_string(),
                provider_id: "provider-1".to_string(),
                name: "key".to_string(),
                auth_type: "bearer".to_string(),
                is_active: true,
                api_formats: None,
                auth_type_by_format: None,
                allow_auth_channel_mismatch_formats: None,
                allowed_models: None,
                capabilities: None,
                rate_multipliers: None,
                global_priority_by_format: None,
                expires_at_unix_secs: None,
                proxy: None,
                fingerprint: None,
                upstream_metadata: None,
                decrypted_api_key: "secret".to_string(),
                decrypted_auth_config: None,
            },
        }
    }

    #[test]
    fn codex_openai_image_url_uses_images_surface() {
        let url = build_openai_image_upstream_url(
            &sample_transport(),
            Some("/v1/images/generations"),
            Some("trace=1"),
        );

        assert_eq!(url, "https://api.openai.com/v1/images/generations?trace=1");
    }

    #[test]
    fn codex_openai_image_edit_url_uses_images_edit_surface() {
        let url = build_openai_image_upstream_url(
            &sample_transport(),
            Some("/v1/images/edits"),
            Some("trace=1"),
        );

        assert_eq!(url, "https://api.openai.com/v1/images/edits?trace=1");
    }

    #[test]
    fn standard_openai_image_url_uses_images_surface() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "openai".to_string();

        let url = build_openai_image_upstream_url(
            &transport,
            Some("/v1/images/generations"),
            Some("trace=1"),
        );

        assert_eq!(url, "https://api.openai.com/v1/images/generations?trace=1");
    }

    #[test]
    fn standard_openai_image_url_preserves_edit_surface() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "openai".to_string();

        let url =
            build_openai_image_upstream_url(&transport, Some("/v1/images/edits"), Some("trace=1"));

        assert_eq!(url, "https://api.openai.com/v1/images/edits?trace=1");
    }

    #[test]
    fn chatgpt_web_is_supported_by_dedicated_openai_image_transport_policy() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "chatgpt_web".to_string();

        assert_eq!(
            openai_image_transport_unsupported_reason(&transport, "openai:image"),
            None
        );
    }

    #[test]
    fn grok_is_supported_by_dedicated_openai_image_transport_policy() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "grok".to_string();

        assert_eq!(
            openai_image_transport_unsupported_reason(&transport, "openai:image"),
            None
        );
    }

    #[test]
    fn codex_is_supported_by_dedicated_openai_image_transport_policy() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "codex".to_string();
        transport.key.auth_type = "oauth".to_string();
        transport.key.decrypted_auth_config = Some(json!({"access_token":"token"}).to_string());

        assert_eq!(
            openai_image_transport_unsupported_reason(&transport, "openai:image"),
            None
        );
    }

    #[test]
    fn grok_oauth_session_is_supported_by_dedicated_openai_image_transport_policy() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "grok".to_string();
        transport.key.auth_type = "oauth".to_string();
        transport.key.decrypted_api_key = String::new();
        transport.key.decrypted_auth_config = Some(json!({"sso_token":"abc"}).to_string());

        assert_eq!(
            openai_image_transport_unsupported_reason(&transport, "openai:image"),
            None
        );
    }

    #[test]
    fn chatgpt_web_oauth_is_supported_by_dedicated_openai_image_transport_policy() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "chatgpt_web".to_string();
        transport.key.auth_type = "oauth".to_string();
        transport.key.decrypted_api_key = String::new();
        transport.key.decrypted_auth_config = Some(json!({"access_token":"token"}).to_string());

        assert_eq!(
            openai_image_transport_unsupported_reason(&transport, "openai:image"),
            None
        );
    }

    #[test]
    fn builds_json_eventstream_headers_and_applies_rules() {
        let transport = sample_transport();
        let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {
            transport: &transport,
            headers: &HeaderMap::new(),
            auth_header: "authorization",
            auth_value: "Bearer secret",
            accept: Some("text/event-stream"),
            header_rules: Some(&json!([
                {"action":"set","key":"x-image-route","value":"codex"}
            ])),
            provider_request_body: &json!({"model":"gpt-5.4-mini"}),
            original_request_body: &json!({"prompt":"draw"}),
        })
        .expect("headers should build");

        assert_eq!(
            headers.get("authorization"),
            Some(&"Bearer secret".to_string())
        );
        assert_eq!(
            headers.get("content-type"),
            Some(&"application/json".to_string())
        );
        assert_eq!(
            headers.get("accept"),
            Some(&"text/event-stream".to_string())
        );
        assert_eq!(headers.get("x-image-route"), Some(&"codex".to_string()));
    }

    #[test]
    fn auth_config_headers_override_default_authorization() {
        let mut transport = sample_transport();
        transport.key.decrypted_auth_config = Some(
            json!({
                "refresh_token": "rt-1",
                "headers": {
                    "authorization": "Bearer imported-session",
                    "chatgpt-account-id": "acct-1"
                }
            })
            .to_string(),
        );

        let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {
            transport: &transport,
            headers: &HeaderMap::new(),
            auth_header: "authorization",
            auth_value: "Bearer refreshed-access-token",
            accept: Some("text/event-stream"),
            header_rules: None,
            provider_request_body: &json!({"model":"gpt-5.4-mini"}),
            original_request_body: &json!({"prompt":"draw"}),
        })
        .expect("headers should build");

        assert_eq!(
            headers.get("authorization"),
            Some(&"Bearer imported-session".to_string())
        );
        assert_eq!(
            headers.get("chatgpt-account-id"),
            Some(&"acct-1".to_string())
        );
    }

    #[test]
    fn standard_openai_compatible_image_headers_can_request_json() {
        let transport = sample_transport();
        let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {
            transport: &transport,
            headers: &HeaderMap::new(),
            auth_header: "authorization",
            auth_value: "Bearer secret",
            accept: Some("application/json"),
            header_rules: None,
            provider_request_body: &json!({
                "model": "upstream-image-model",
                "prompt": "draw a city",
            }),
            original_request_body: &json!({"prompt":"draw a city"}),
        })
        .expect("headers should build");

        assert_eq!(headers.get("accept"), Some(&"application/json".to_string()));
        assert_eq!(
            headers.get("content-type"),
            Some(&"application/json".to_string())
        );
    }

    #[test]
    fn codex_images_omits_explicit_accept_header() {
        let transport = sample_transport();
        let mut request_headers = HeaderMap::new();
        request_headers.insert("accept", "text/event-stream".parse().expect("valid header"));
        let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {
            transport: &transport,
            headers: &request_headers,
            auth_header: "authorization",
            auth_value: "Bearer secret",
            accept: None,
            header_rules: Some(&json!([
                {"action":"set","key":"accept","value":"application/json"}
            ])),
            provider_request_body: &json!({
                "model": "gpt-image-2",
                "prompt": "draw a city",
            }),
            original_request_body: &json!({"prompt":"draw a city"}),
        })
        .expect("headers should build");

        assert!(!headers.contains_key("accept"));
    }

    #[test]
    fn standard_openai_compatible_image_url_supports_aether_api_root() {
        let mut transport = sample_transport();
        transport.provider.provider_type = "custom".to_string();
        transport.endpoint.base_url = "https://upstream-aether.example/v1".to_string();

        let url = build_openai_image_upstream_url(
            &transport,
            Some("/v1/images/generations"),
            Some("trace=1"),
        );

        assert_eq!(
            url,
            "https://upstream-aether.example/v1/images/generations?trace=1"
        );
    }
}