Aether/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs

use std::collections::BTreeMap;

use super::{apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers};
use crate::ai_serving::planner::standard::build_local_openai_responses_request_body;
use http::{HeaderMap, HeaderValue};
use serde_json::json;

#[test]
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
    let mut body = json!({
        "model": "gpt-5.4",
        "max_output_tokens": 128,
        "temperature": 0.3,
        "top_p": 0.9,
        "metadata": {"client": "desktop"},
        "store": true
    });

    apply_codex_openai_responses_special_body_edits(
        &mut body,
        "codex",
        "openai:responses",
        None,
        None,
    );

    assert!(body.get("max_output_tokens").is_none());
    assert!(body.get("temperature").is_none());
    assert!(body.get("top_p").is_none());
    assert!(body.get("metadata").is_none());
    assert_eq!(body["store"], false);
    assert!(body.get("instructions").is_none());
    assert_eq!(body["include"], json!(["reasoning.encrypted_content"]));
    assert_eq!(body["parallel_tool_calls"], true);
    assert_eq!(body["reasoning"]["effort"], "medium");
    assert!(body["reasoning"].get("summary").is_none());
}

#[test]
fn local_openai_responses_codex_body_wraps_string_input_for_backend() {
    let body = json!({
        "model": "gpt-5",
        "input": "hello"
    });

    let provider_request_body = build_local_openai_responses_request_body(
        &body,
        "gpt-5-upstream",
        false,
        false,
        "codex",
        "openai:responses",
        None,
        Some("key-123"),
        &HeaderMap::new(),
        false,
    )
    .expect("codex local openai responses body should build");

    assert_eq!(
        provider_request_body["input"],
        json!([{
            "type": "message",
            "role": "user",
            "content": [{
                "type": "input_text",
                "text": "hello"
            }]
        }])
    );
}

#[test]
fn strips_store_for_compact_even_when_body_rules_handle_it() {
    let body_rules = json!([
        {"action":"set","path":"store","value":true},
        {"action":"set","path":"instructions","value":"Keep custom"},
        {"action":"set","path":"metadata","value":{"client":"desktop","mode":"custom"}},
        {"action":"set","path":"top_p","value":0.5}
    ]);
    let mut body = json!({
        "model": "gpt-5.4",
        "max_output_tokens": 128,
        "metadata": {"client": "desktop", "mode": "custom"},
        "store": true,
        "instructions": "Keep custom",
        "top_p": 0.5
    });

    apply_codex_openai_responses_special_body_edits(
        &mut body,
        "codex",
        "openai:responses:compact",
        Some(&body_rules),
        None,
    );

    assert!(body.get("max_output_tokens").is_none());
    assert!(body.get("store").is_none());
    assert_eq!(body["instructions"], "Keep custom");
    assert!(body.get("metadata").is_none());
    assert!(body.get("top_p").is_none());
    assert_eq!(body["parallel_tool_calls"], true);
    assert!(body.as_object().is_some_and(|object| {
        object.keys().all(|field| {
            matches!(
                field.as_str(),
                "model"
                    | "input"
                    | "instructions"
                    | "tools"
                    | "parallel_tool_calls"
                    | "reasoning"
                    | "service_tier"
                    | "prompt_cache_key"
                    | "text"
            )
        })
    }));
}

#[test]
fn does_not_synthesize_prompt_cache_key_from_api_key_identity() {
    let mut body = json!({
        "model": "gpt-5",
        "input": "hello",
    });

    apply_codex_openai_responses_special_body_edits(
        &mut body,
        "codex",
        "openai:responses",
        None,
        Some("key-123"),
    );

    assert!(body.get("prompt_cache_key").is_none());
}

#[test]
fn keeps_existing_prompt_cache_key_for_codex_requests() {
    let mut body = json!({
        "model": "gpt-5",
        "input": "hello",
        "prompt_cache_key": "existing-key",
    });

    apply_codex_openai_responses_special_body_edits(
        &mut body,
        "codex",
        "openai:responses",
        None,
        Some("key-123"),
    );

    assert_eq!(body["prompt_cache_key"], "existing-key");
}

#[test]
fn injects_identity_headers_without_deriving_session_headers_from_body() {
    let mut headers = BTreeMap::new();
    let body = json!({
        "model": "gpt-5",
        "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
    });

    apply_codex_openai_special_headers(
        &mut headers,
        &body,
        &HeaderMap::new(),
        "codex",
        "openai:responses",
        Some("trace-codex-123"),
        Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
    );

    assert_eq!(
        headers.get("chatgpt-account-id"),
        Some(&"acc-123".to_string())
    );
    assert_eq!(headers.get("x-client-request-id"), None);
    assert_eq!(
        headers.get("user-agent"),
        Some(&"codex_cli_rs/0.144.1".to_string())
    );
    assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
    assert!(!headers.contains_key("version"));
    assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string()));
    assert_eq!(headers.get("session-id"), None);
    assert_eq!(headers.get("thread-id"), None);
}

#[test]
fn injects_only_codex_client_headers_for_images_requests() {
    let mut headers = BTreeMap::new();
    apply_codex_openai_special_headers(
        &mut headers,
        &json!({
            "model": "gpt-image-2",
            "prompt": "draw a city"
        }),
        &HeaderMap::new(),
        "codex",
        "openai:image",
        Some("trace-codex-image-123"),
        Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
    );

    assert_eq!(
        headers.get("chatgpt-account-id"),
        Some(&"acc-123".to_string())
    );
    assert_eq!(
        headers.get("user-agent"),
        Some(&"codex_cli_rs/0.144.1".to_string())
    );
    assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
    assert!(!headers.contains_key("version"));
    assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string()));
    for name in ["x-client-request-id", "session-id", "thread-id"] {
        assert!(
            !headers.contains_key(name),
            "unexpected Images header: {name}"
        );
    }
}

#[test]
fn preserves_client_context_headers_and_enforces_codex_auth_identity_headers() {
    let mut headers = BTreeMap::new();
    headers.insert(
        "x-client-request-id".to_string(),
        "kept-by-rule-request".to_string(),
    );
    headers.insert("session-id".to_string(), "kept-by-rule-session".to_string());
    headers.insert("thread-id".to_string(), "kept-by-rule-thread".to_string());
    headers.insert(
        "chatgpt-account-id".to_string(),
        "configured-spoof".to_string(),
    );
    headers.insert(
        "x-openai-fedramp".to_string(),
        "configured-false".to_string(),
    );
    let body = json!({
        "model": "gpt-5",
        "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
    });
    let mut original_headers = HeaderMap::new();
    original_headers.insert(
        "x-client-request-id",
        HeaderValue::from_static("user-specified-request"),
    );
    original_headers.insert(
        "session-id",
        HeaderValue::from_static("user-specified-session"),
    );
    original_headers.insert(
        "thread-id",
        HeaderValue::from_static("user-specified-thread"),
    );
    original_headers.insert(
        "user-agent",
        HeaderValue::from_static("user-specified-agent"),
    );
    original_headers.insert(
        "originator",
        HeaderValue::from_static("user-specified-originator"),
    );
    original_headers.insert("version", HeaderValue::from_static("user-version"));
    original_headers.insert("x-openai-fedramp", HeaderValue::from_static("user-fedramp"));
    original_headers.insert(
        "chatgpt-account-id",
        HeaderValue::from_static("user-account"),
    );

    apply_codex_openai_special_headers(
        &mut headers,
        &body,
        &original_headers,
        "codex",
        "openai:responses",
        Some("trace-codex-123"),
        Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
    );

    assert_eq!(
        headers.get("x-client-request-id"),
        Some(&"kept-by-rule-request".to_string())
    );
    assert!(!headers.contains_key("user-agent"));
    assert!(!headers.contains_key("originator"));
    assert!(!headers.contains_key("version"));
    assert_eq!(
        headers.get("chatgpt-account-id"),
        Some(&"acc-123".to_string())
    );
    assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string()));
    assert_eq!(
        headers.get("session-id"),
        Some(&"kept-by-rule-session".to_string())
    );
    assert_eq!(
        headers.get("thread-id"),
        Some(&"kept-by-rule-thread".to_string())
    );
}

#[test]
fn compact_does_not_derive_session_headers_from_body() {
    let mut headers = BTreeMap::new();
    let body = json!({
        "model": "gpt-5",
        "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
    });

    apply_codex_openai_special_headers(
        &mut headers,
        &body,
        &HeaderMap::new(),
        "codex",
        "openai:responses:compact",
        Some("trace-codex-compact-123"),
        Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
    );

    assert_eq!(
        headers.get("chatgpt-account-id"),
        Some(&"acc-123".to_string())
    );
    assert_eq!(headers.get("x-client-request-id"), None);
    assert_eq!(
        headers.get("user-agent"),
        Some(&"codex_cli_rs/0.144.1".to_string())
    );
    assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
    assert!(!headers.contains_key("version"));
    assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string()));
    assert_eq!(headers.get("session-id"), None);
    assert_eq!(headers.get("thread-id"), None);
}