codex/codex-rs/app-server-protocol/src/export.rs

use crate::ClientNotification;
use crate::ClientRequest;
use crate::ServerNotification;
use crate::ServerRequest;
use crate::export_client_param_schemas;
use crate::export_client_response_schemas;
use crate::export_client_responses;
use crate::export_server_param_schemas;
use crate::export_server_response_schemas;
use crate::export_server_responses;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use schemars::JsonSchema;
use schemars::schema::RootSchema;
use schemars::schema_for;
use serde::Serialize;
use serde_json::Map;
use serde_json::Value;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::ffi::OsStr;
use std::fs;
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use ts_rs::TS;

const HEADER: &str = "// GENERATED CODE! DO NOT MODIFY BY HAND!\n\n";
type JsonSchemaEmitter = fn(&Path) -> Result<RootSchema>;

pub fn generate_types(out_dir: &Path, prettier: Option<&Path>) -> Result<()> {
    generate_ts(out_dir, prettier)?;
    generate_json(out_dir)?;
    Ok(())
}

pub fn generate_ts(out_dir: &Path, prettier: Option<&Path>) -> Result<()> {
    let v2_out_dir = out_dir.join("v2");
    ensure_dir(out_dir)?;
    ensure_dir(&v2_out_dir)?;

    ClientRequest::export_all_to(out_dir)?;
    export_client_responses(out_dir)?;
    ClientNotification::export_all_to(out_dir)?;

    ServerRequest::export_all_to(out_dir)?;
    export_server_responses(out_dir)?;
    ServerNotification::export_all_to(out_dir)?;

    generate_index_ts(out_dir)?;
    generate_index_ts(&v2_out_dir)?;

    // Ensure our header is present on all TS files (root + subdirs like v2/).
    let ts_files = ts_files_in_recursive(out_dir)?;
    for file in &ts_files {
        prepend_header_if_missing(file)?;
    }

    // Optionally run Prettier on all generated TS files.
    if let Some(prettier_bin) = prettier
        && !ts_files.is_empty()
    {
        let status = Command::new(prettier_bin)
            .arg("--write")
            .args(ts_files.iter().map(|p| p.as_os_str()))
            .status()
            .with_context(|| format!("Failed to invoke Prettier at {}", prettier_bin.display()))?;
        if !status.success() {
            return Err(anyhow!("Prettier failed with status {status}"));
        }
    }

    Ok(())
}

pub fn generate_json(out_dir: &Path) -> Result<()> {
    ensure_dir(out_dir)?;
    // Emit schemas for core envelope/wire types that may not be referenced
    // by params/responses but are useful as standalone:
    let mut bundle: BTreeMap<String, RootSchema> = BTreeMap::new();

    let envelope_emitters: &[(&str, JsonSchemaEmitter)] = &[
        ("RequestId", |d| {
            write_json_schema_with_return::<crate::RequestId>(d, "RequestId")
        }),
        ("JSONRPCMessage", |d| {
            write_json_schema_with_return::<crate::JSONRPCMessage>(d, "JSONRPCMessage")
        }),
        ("JSONRPCRequest", |d| {
            write_json_schema_with_return::<crate::JSONRPCRequest>(d, "JSONRPCRequest")
        }),
        ("JSONRPCNotification", |d| {
            write_json_schema_with_return::<crate::JSONRPCNotification>(d, "JSONRPCNotification")
        }),
        ("JSONRPCResponse", |d| {
            write_json_schema_with_return::<crate::JSONRPCResponse>(d, "JSONRPCResponse")
        }),
        ("JSONRPCError", |d| {
            write_json_schema_with_return::<crate::JSONRPCError>(d, "JSONRPCError")
        }),
        ("JSONRPCErrorError", |d| {
            write_json_schema_with_return::<crate::JSONRPCErrorError>(d, "JSONRPCErrorError")
        }),
        ("ClientRequest", |d| {
            write_json_schema_with_return::<crate::ClientRequest>(d, "ClientRequest")
        }),
        ("ServerRequest", |d| {
            write_json_schema_with_return::<crate::ServerRequest>(d, "ServerRequest")
        }),
        ("ClientNotification", |d| {
            write_json_schema_with_return::<crate::ClientNotification>(d, "ClientNotification")
        }),
        ("ServerNotification", |d| {
            write_json_schema_with_return::<crate::ServerNotification>(d, "ServerNotification")
        }),
    ];
    for (name, emit) in envelope_emitters {
        let schema = emit(out_dir)?;
        bundle.insert((*name).to_string(), schema);
    }

    // Have the macros generate per-type JSON for params and responses.
    export_client_param_schemas(out_dir)?;
    export_client_response_schemas(out_dir)?;
    export_server_param_schemas(out_dir)?;
    export_server_response_schemas(out_dir)?;

    let mut definitions = Map::new();

    const SPECIAL_DEFINITIONS: &[&str] = &[
        "ClientNotification",
        "ClientRequest",
        "EventMsg",
        "FileChange",
        "InputItem",
        "ParsedCommand",
        "SandboxPolicy",
        "ServerNotification",
        "ServerRequest",
    ];

    // Merge all generated per-type JSON files (including the envelopes above)
    // into a single definitions bundle.
    for entry in fs::read_dir(out_dir)
        .with_context(|| format!("Failed to read dir {}", out_dir.display()))?
    {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            continue;
        }
        if path.extension() != Some(OsStr::new("json")) {
            continue;
        }
        if path.file_name().and_then(OsStr::to_str)
            == Some("codex_app_server_protocol.schemas.json")
        {
            // Skip the bundle we’re about to (re)generate.
            continue;
        }

        let name = path
            .file_stem()
            .and_then(OsStr::to_str)
            .map(str::to_owned)
            .ok_or_else(|| anyhow!(format!("Invalid schema file name {}", path.display())))?;

        let mut schema_value: Value = serde_json::from_str(
            &fs::read_to_string(&path)
                .with_context(|| format!("Failed to read {}", path.display()))?,
        )
        .with_context(|| format!("Failed to parse JSON from {}", path.display()))?;

        // Ensure the same normalization logic applies when building the bundle.
        annotate_schema(&mut schema_value, Some(name.as_str()));

        if let Value::Object(ref mut obj) = schema_value
            && let Some(defs) = obj.remove("definitions")
            && let Value::Object(defs_obj) = defs
        {
            for (def_name, mut def_schema) in defs_obj {
                if !SPECIAL_DEFINITIONS.contains(&def_name.as_str()) {
                    annotate_schema(&mut def_schema, Some(def_name.as_str()));
                    definitions.insert(def_name, def_schema);
                }
            }
        }

        definitions.insert(name, schema_value);
    }

    let mut root = Map::new();
    root.insert(
        "$schema".to_string(),
        Value::String("http://json-schema.org/draft-07/schema#".into()),
    );
    root.insert(
        "title".to_string(),
        Value::String("CodexAppServerProtocol".into()),
    );
    root.insert("type".to_string(), Value::String("object".into()));
    root.insert("definitions".to_string(), Value::Object(definitions));

    write_pretty_json(
        out_dir.join("codex_app_server_protocol.schemas.json"),
        &Value::Object(root),
    )?;

    Ok(())
}

fn write_json_schema_with_return<T>(out_dir: &Path, name: &str) -> Result<RootSchema>
where
    T: JsonSchema,
{
    let file_stem = name.trim();
    let schema = schema_for!(T);
    let mut schema_value = serde_json::to_value(schema)?;
    annotate_schema(&mut schema_value, Some(file_stem));
    write_pretty_json(out_dir.join(format!("{file_stem}.json")), &schema_value)
        .with_context(|| format!("Failed to write JSON schema for {file_stem}"))?;
    let annotated_schema = serde_json::from_value(schema_value)?;
    Ok(annotated_schema)
}

pub(crate) fn write_json_schema<T>(out_dir: &Path, name: &str) -> Result<()>
where
    T: JsonSchema,
{
    write_json_schema_with_return::<T>(out_dir, name).map(|_| ())
}

fn write_pretty_json(path: PathBuf, value: &impl Serialize) -> Result<()> {
    let json = serde_json::to_vec_pretty(value)
        .with_context(|| format!("Failed to serialize JSON schema to {}", path.display()))?;
    fs::write(&path, json).with_context(|| format!("Failed to write {}", path.display()))?;
    Ok(())
}

fn variant_definition_name(base: &str, variant: &Value) -> Option<String> {
    if let Some(props) = variant.get("properties").and_then(Value::as_object) {
        if let Some(method_literal) = literal_from_property(props, "method") {
            let pascal = to_pascal_case(method_literal);
            return Some(match base {
                "ClientRequest" | "ServerRequest" => format!("{pascal}Request"),
                "ClientNotification" | "ServerNotification" => format!("{pascal}Notification"),
                _ => format!("{pascal}{base}"),
            });
        }

        if let Some(type_literal) = literal_from_property(props, "type") {
            let pascal = to_pascal_case(type_literal);
            return Some(match base {
                "EventMsg" => format!("{pascal}EventMsg"),
                _ => format!("{pascal}{base}"),
            });
        }

        if let Some(mode_literal) = literal_from_property(props, "mode") {
            let pascal = to_pascal_case(mode_literal);
            return Some(match base {
                "SandboxPolicy" => format!("{pascal}SandboxPolicy"),
                _ => format!("{pascal}{base}"),
            });
        }

        if props.len() == 1
            && let Some(key) = props.keys().next()
        {
            let pascal = to_pascal_case(key);
            return Some(format!("{pascal}{base}"));
        }
    }

    if let Some(required) = variant.get("required").and_then(Value::as_array)
        && required.len() == 1
        && let Some(key) = required[0].as_str()
    {
        let pascal = to_pascal_case(key);
        return Some(format!("{pascal}{base}"));
    }

    None
}

fn literal_from_property<'a>(props: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
    props.get(key).and_then(string_literal)
}

fn string_literal(value: &Value) -> Option<&str> {
    value.get("const").and_then(Value::as_str).or_else(|| {
        value
            .get("enum")
            .and_then(Value::as_array)
            .and_then(|arr| arr.first())
            .and_then(Value::as_str)
    })
}

fn annotate_schema(value: &mut Value, base: Option<&str>) {
    match value {
        Value::Object(map) => annotate_object(map, base),
        Value::Array(items) => {
            for item in items {
                annotate_schema(item, base);
            }
        }
        _ => {}
    }
}

fn annotate_object(map: &mut Map<String, Value>, base: Option<&str>) {
    let owner = map.get("title").and_then(Value::as_str).map(str::to_owned);
    if let Some(owner) = owner.as_deref()
        && let Some(Value::Object(props)) = map.get_mut("properties")
    {
        set_discriminator_titles(props, owner);
    }

    if let Some(Value::Array(variants)) = map.get_mut("oneOf") {
        annotate_variant_list(variants, base);
    }
    if let Some(Value::Array(variants)) = map.get_mut("anyOf") {
        annotate_variant_list(variants, base);
    }

    if let Some(Value::Object(defs)) = map.get_mut("definitions") {
        for (name, schema) in defs.iter_mut() {
            annotate_schema(schema, Some(name.as_str()));
        }
    }

    if let Some(Value::Object(defs)) = map.get_mut("$defs") {
        for (name, schema) in defs.iter_mut() {
            annotate_schema(schema, Some(name.as_str()));
        }
    }

    if let Some(Value::Object(props)) = map.get_mut("properties") {
        for value in props.values_mut() {
            annotate_schema(value, base);
        }
    }

    if let Some(items) = map.get_mut("items") {
        annotate_schema(items, base);
    }

    if let Some(additional) = map.get_mut("additionalProperties") {
        annotate_schema(additional, base);
    }

    for (key, child) in map.iter_mut() {
        match key.as_str() {
            "oneOf"
            | "anyOf"
            | "definitions"
            | "$defs"
            | "properties"
            | "items"
            | "additionalProperties" => {}
            _ => annotate_schema(child, base),
        }
    }
}

fn annotate_variant_list(variants: &mut [Value], base: Option<&str>) {
    let mut seen = HashSet::new();

    for variant in variants.iter() {
        if let Some(name) = variant_title(variant) {
            seen.insert(name.to_owned());
        }
    }

    for variant in variants.iter_mut() {
        let mut variant_name = variant_title(variant).map(str::to_owned);

        if variant_name.is_none()
            && let Some(base_name) = base
            && let Some(name) = variant_definition_name(base_name, variant)
        {
            let mut candidate = name.clone();
            let mut index = 2;
            while seen.contains(&candidate) {
                candidate = format!("{name}{index}");
                index += 1;
            }
            if let Some(obj) = variant.as_object_mut() {
                obj.insert("title".into(), Value::String(candidate.clone()));
            }
            seen.insert(candidate.clone());
            variant_name = Some(candidate);
        }

        if let Some(name) = variant_name.as_deref()
            && let Some(obj) = variant.as_object_mut()
            && let Some(Value::Object(props)) = obj.get_mut("properties")
        {
            set_discriminator_titles(props, name);
        }

        annotate_schema(variant, base);
    }
}

const DISCRIMINATOR_KEYS: &[&str] = &["type", "method", "mode", "status", "role", "reason"];

fn set_discriminator_titles(props: &mut Map<String, Value>, owner: &str) {
    for key in DISCRIMINATOR_KEYS {
        if let Some(prop_schema) = props.get_mut(*key)
            && string_literal(prop_schema).is_some()
            && let Value::Object(prop_obj) = prop_schema
        {
            if prop_obj.contains_key("title") {
                continue;
            }
            let suffix = to_pascal_case(key);
            prop_obj.insert("title".into(), Value::String(format!("{owner}{suffix}")));
        }
    }
}

fn variant_title(value: &Value) -> Option<&str> {
    value
        .as_object()
        .and_then(|obj| obj.get("title"))
        .and_then(Value::as_str)
}

fn to_pascal_case(input: &str) -> String {
    let mut result = String::new();
    let mut capitalize_next = true;

    for c in input.chars() {
        if c == '_' || c == '-' {
            capitalize_next = true;
            continue;
        }

        if capitalize_next {
            result.extend(c.to_uppercase());
            capitalize_next = false;
        } else {
            result.push(c);
        }
    }

    result
}

fn ensure_dir(dir: &Path) -> Result<()> {
    fs::create_dir_all(dir)
        .with_context(|| format!("Failed to create output directory {}", dir.display()))
}

fn prepend_header_if_missing(path: &Path) -> Result<()> {
    let mut content = String::new();
    {
        let mut f = fs::File::open(path)
            .with_context(|| format!("Failed to open {} for reading", path.display()))?;
        f.read_to_string(&mut content)
            .with_context(|| format!("Failed to read {}", path.display()))?;
    }

    if content.starts_with(HEADER) {
        return Ok(());
    }

    let mut f = fs::File::create(path)
        .with_context(|| format!("Failed to open {} for writing", path.display()))?;
    f.write_all(HEADER.as_bytes())
        .with_context(|| format!("Failed to write header to {}", path.display()))?;
    f.write_all(content.as_bytes())
        .with_context(|| format!("Failed to write content to {}", path.display()))?;
    Ok(())
}

fn ts_files_in(dir: &Path) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();
    for entry in
        fs::read_dir(dir).with_context(|| format!("Failed to read dir {}", dir.display()))?
    {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() && path.extension() == Some(OsStr::new("ts")) {
            files.push(path);
        }
    }
    files.sort();
    Ok(files)
}

fn ts_files_in_recursive(dir: &Path) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();
    let mut stack = vec![dir.to_path_buf()];
    while let Some(d) = stack.pop() {
        for entry in
            fs::read_dir(&d).with_context(|| format!("Failed to read dir {}", d.display()))?
        {
            let entry = entry?;
            let path = entry.path();
            if path.is_dir() {
                stack.push(path);
            } else if path.is_file() && path.extension() == Some(OsStr::new("ts")) {
                files.push(path);
            }
        }
    }
    files.sort();
    Ok(files)
}

fn generate_index_ts(out_dir: &Path) -> Result<PathBuf> {
    let mut entries: Vec<String> = Vec::new();
    let mut stems: Vec<String> = ts_files_in(out_dir)?
        .into_iter()
        .filter_map(|p| {
            let stem = p.file_stem()?.to_string_lossy().into_owned();
            if stem == "index" { None } else { Some(stem) }
        })
        .collect();
    stems.sort();
    stems.dedup();

    for name in stems {
        entries.push(format!("export type {{ {name} }} from \"./{name}\";\n"));
    }

    // If this is the root out_dir and a ./v2 folder exists with TS files,
    // expose it as a namespace to avoid symbol collisions at the root.
    let v2_dir = out_dir.join("v2");
    let has_v2_ts = ts_files_in(&v2_dir).map(|v| !v.is_empty()).unwrap_or(false);
    if has_v2_ts {
        entries.push("export * as v2 from \"./v2\";\n".to_string());
    }

    let mut content =
        String::with_capacity(HEADER.len() + entries.iter().map(String::len).sum::<usize>());
    content.push_str(HEADER);
    for line in &entries {
        content.push_str(line);
    }

    let index_path = out_dir.join("index.ts");
    let mut f = fs::File::create(&index_path)
        .with_context(|| format!("Failed to create {}", index_path.display()))?;
    f.write_all(content.as_bytes())
        .with_context(|| format!("Failed to write {}", index_path.display()))?;
    Ok(index_path)
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use std::collections::BTreeSet;
    use std::fs;
    use std::path::PathBuf;
    use uuid::Uuid;

    #[test]
    fn generated_ts_has_no_optional_nullable_fields() -> Result<()> {
        // Assert that there are no types of the form "?: T | null" in the generated TS files.
        let output_dir = std::env::temp_dir().join(format!("codex_ts_types_{}", Uuid::now_v7()));
        fs::create_dir(&output_dir)?;

        struct TempDirGuard(PathBuf);

        impl Drop for TempDirGuard {
            fn drop(&mut self) {
                let _ = fs::remove_dir_all(&self.0);
            }
        }

        let _guard = TempDirGuard(output_dir.clone());

        generate_ts(&output_dir, None)?;

        let mut undefined_offenders = Vec::new();
        let mut optional_nullable_offenders = BTreeSet::new();
        let mut stack = vec![output_dir];
        while let Some(dir) = stack.pop() {
            for entry in fs::read_dir(&dir)? {
                let entry = entry?;
                let path = entry.path();
                if path.is_dir() {
                    stack.push(path);
                    continue;
                }

                if matches!(path.extension().and_then(|ext| ext.to_str()), Some("ts")) {
                    let contents = fs::read_to_string(&path)?;
                    if contents.contains("| undefined") {
                        undefined_offenders.push(path.clone());
                    }

                    const SKIP_PREFIXES: &[&str] = &[
                        "const ",
                        "let ",
                        "var ",
                        "export const ",
                        "export let ",
                        "export var ",
                    ];

                    let mut search_start = 0;
                    while let Some(idx) = contents[search_start..].find("| null") {
                        let abs_idx = search_start + idx;
                        // Find the property-colon for this field by scanning forward
                        // from the start of the segment and ignoring nested braces,
                        // brackets, and parens. This avoids colons inside nested
                        // type literals like `{ [k in string]?: string }`.

                        let line_start_idx =
                            contents[..abs_idx].rfind('\n').map(|i| i + 1).unwrap_or(0);

                        let mut segment_start_idx = line_start_idx;
                        if let Some(rel_idx) = contents[line_start_idx..abs_idx].rfind(',') {
                            segment_start_idx = segment_start_idx.max(line_start_idx + rel_idx + 1);
                        }
                        if let Some(rel_idx) = contents[line_start_idx..abs_idx].rfind('{') {
                            segment_start_idx = segment_start_idx.max(line_start_idx + rel_idx + 1);
                        }
                        if let Some(rel_idx) = contents[line_start_idx..abs_idx].rfind('}') {
                            segment_start_idx = segment_start_idx.max(line_start_idx + rel_idx + 1);
                        }

                        // Scan forward for the colon that separates the field name from its type.
                        let mut level_brace = 0_i32;
                        let mut level_brack = 0_i32;
                        let mut level_paren = 0_i32;
                        let mut in_single = false;
                        let mut in_double = false;
                        let mut escape = false;
                        let mut prop_colon_idx = None;
                        for (i, ch) in contents[segment_start_idx..abs_idx].char_indices() {
                            let idx_abs = segment_start_idx + i;
                            if escape {
                                escape = false;
                                continue;
                            }
                            match ch {
                                '\\' => {
                                    // Only treat as escape when inside a string.
                                    if in_single || in_double {
                                        escape = true;
                                    }
                                }
                                '\'' => {
                                    if !in_double {
                                        in_single = !in_single;
                                    }
                                }
                                '"' => {
                                    if !in_single {
                                        in_double = !in_double;
                                    }
                                }
                                '{' if !in_single && !in_double => level_brace += 1,
                                '}' if !in_single && !in_double => level_brace -= 1,
                                '[' if !in_single && !in_double => level_brack += 1,
                                ']' if !in_single && !in_double => level_brack -= 1,
                                '(' if !in_single && !in_double => level_paren += 1,
                                ')' if !in_single && !in_double => level_paren -= 1,
                                ':' if !in_single
                                    && !in_double
                                    && level_brace == 0
                                    && level_brack == 0
                                    && level_paren == 0 =>
                                {
                                    prop_colon_idx = Some(idx_abs);
                                    break;
                                }
                                _ => {}
                            }
                        }

                        let Some(colon_idx) = prop_colon_idx else {
                            search_start = abs_idx + 5;
                            continue;
                        };

                        let mut field_prefix = contents[segment_start_idx..colon_idx].trim();
                        if field_prefix.is_empty() {
                            search_start = abs_idx + 5;
                            continue;
                        }

                        if let Some(comment_idx) = field_prefix.rfind("*/") {
                            field_prefix = field_prefix[comment_idx + 2..].trim_start();
                        }

                        if field_prefix.is_empty() {
                            search_start = abs_idx + 5;
                            continue;
                        }

                        if SKIP_PREFIXES
                            .iter()
                            .any(|prefix| field_prefix.starts_with(prefix))
                        {
                            search_start = abs_idx + 5;
                            continue;
                        }

                        if field_prefix.contains('(') {
                            search_start = abs_idx + 5;
                            continue;
                        }

                        // If the last non-whitespace before ':' is '?', then this is an
                        // optional field with a nullable type (i.e., "?: T | null"),
                        // which we explicitly disallow.
                        if field_prefix.chars().rev().find(|c| !c.is_whitespace()) == Some('?') {
                            let line_number =
                                contents[..abs_idx].chars().filter(|c| *c == '\n').count() + 1;
                            let offending_line_end = contents[line_start_idx..]
                                .find('\n')
                                .map(|i| line_start_idx + i)
                                .unwrap_or(contents.len());
                            let offending_snippet =
                                contents[line_start_idx..offending_line_end].trim();

                            optional_nullable_offenders.insert(format!(
                                "{}:{}: {offending_snippet}",
                                path.display(),
                                line_number
                            ));
                        }

                        search_start = abs_idx + 5;
                    }
                }
            }
        }

        assert!(
            undefined_offenders.is_empty(),
            "Generated TypeScript still includes unions with `undefined` in {undefined_offenders:?}"
        );

        // If this assertion fails, it means a field was generated as
        // "?: T | null" — i.e., both optional (undefined) and nullable (null).
        // We only want either "?: T" or ": T | null".
        assert!(
            optional_nullable_offenders.is_empty(),
            "Generated TypeScript has optional fields with nullable types (disallowed '?: T | null'), add #[ts(optional)] to fix:\n{optional_nullable_offenders:?}"
        );

        Ok(())
    }
}