codex/codex-rs/tools/src/code_mode.rs

use crate::ResponsesApiNamespaceTool;
use crate::ToolName;
use crate::ToolOutputSchema;
use crate::ToolSpec;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ToolDefinition as CodeModeToolDefinition;

/// Augment tool descriptions with code-mode-specific exec samples.
pub fn augment_tool_spec_for_code_mode(
    spec: ToolSpec,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> ToolSpec {
    match spec {
        ToolSpec::Function(mut tool) => {
            let Some(description) = augmented_description_for_spec(
                &ToolSpec::Function(tool.clone()),
                code_mode_input_schema_max_bytes,
            ) else {
                return ToolSpec::Function(tool);
            };
            tool.description = description;
            ToolSpec::Function(tool)
        }
        ToolSpec::Freeform(mut tool) => {
            let Some(description) = augmented_description_for_spec(
                &ToolSpec::Freeform(tool.clone()),
                code_mode_input_schema_max_bytes,
            ) else {
                return ToolSpec::Freeform(tool);
            };
            tool.description = description;
            ToolSpec::Freeform(tool)
        }
        ToolSpec::Namespace(mut namespace) => {
            for tool in &mut namespace.tools {
                match tool {
                    ResponsesApiNamespaceTool::Function(tool) => {
                        let tool_name =
                            ToolName::namespaced(namespace.name.clone(), tool.name.clone());
                        let definition = CodeModeToolDefinition {
                            name: code_mode_name_for_tool_name(&tool_name),
                            tool_name,
                            description: tool.description.clone(),
                            kind: CodeModeToolKind::Function,
                            input_schema: serde_json::to_value(&tool.parameters).ok(),
                            input_schema_max_bytes: Some(effective_input_schema_max_bytes(
                                tool.parameters.mcp_input_schema_max_bytes,
                                code_mode_input_schema_max_bytes,
                            )),
                            output_schema: tool
                                .output_schema
                                .as_ref()
                                .map(ToolOutputSchema::to_value),
                        };
                        tool.description =
                            codex_code_mode::augment_tool_definition(definition).description;
                    }
                    ResponsesApiNamespaceTool::Custom(tool) => {
                        let tool_name =
                            ToolName::namespaced(namespace.name.clone(), tool.name.clone());
                        let definition = CodeModeToolDefinition {
                            name: code_mode_name_for_tool_name(&tool_name),
                            tool_name,
                            description: tool.description.clone(),
                            kind: CodeModeToolKind::Freeform,
                            input_schema: None,
                            input_schema_max_bytes: None,
                            output_schema: None,
                        };
                        tool.description =
                            codex_code_mode::augment_tool_definition(definition).description;
                    }
                }
            }
            ToolSpec::Namespace(namespace)
        }
        other => other,
    }
}

/// Convert a supported nested tool spec into the code-mode runtime shape,
/// including the code-mode-specific description sample.
pub fn tool_spec_to_code_mode_tool_definition(spec: &ToolSpec) -> Option<CodeModeToolDefinition> {
    let definition =
        code_mode_tool_definition_for_spec(spec, /*code_mode_input_schema_max_bytes*/ None)?;
    codex_code_mode::is_code_mode_nested_tool(&definition.name)
        .then(|| codex_code_mode::augment_tool_definition(definition))
}

pub fn collect_code_mode_tool_definitions<'a>(
    specs: impl IntoIterator<Item = &'a ToolSpec>,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> Vec<CodeModeToolDefinition> {
    let mut tool_definitions = specs
        .into_iter()
        .flat_map(|spec| {
            let mut definitions =
                code_mode_tool_definitions_for_spec(spec, code_mode_input_schema_max_bytes);
            if let ToolSpec::Namespace(namespace) = spec {
                let namespace_description = namespace.description.trim();
                if !namespace_description.is_empty() {
                    for definition in &mut definitions {
                        definition.description =
                            format!("{namespace_description}\n\n{}", definition.description);
                    }
                }
            }
            definitions
        })
        .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name))
        .map(codex_code_mode::augment_tool_definition)
        .collect::<Vec<_>>();
    tool_definitions.sort_by(|left, right| left.name.cmp(&right.name));
    tool_definitions.dedup_by(|left, right| left.name == right.name);
    tool_definitions
}

pub fn collect_code_mode_exec_prompt_tool_definitions<'a>(
    specs: impl IntoIterator<Item = &'a ToolSpec>,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> Vec<CodeModeToolDefinition> {
    let mut tool_definitions = specs
        .into_iter()
        .flat_map(|spec| {
            code_mode_tool_definitions_for_spec(spec, code_mode_input_schema_max_bytes)
        })
        .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name))
        .collect::<Vec<_>>();
    tool_definitions.sort_by(|left, right| left.name.cmp(&right.name));
    tool_definitions.dedup_by(|left, right| left.name == right.name);
    tool_definitions
}

fn augmented_description_for_spec(
    spec: &ToolSpec,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> Option<String> {
    code_mode_tool_definition_for_spec(spec, code_mode_input_schema_max_bytes)
        .map(codex_code_mode::augment_tool_definition)
        .map(|definition| definition.description)
}

fn code_mode_tool_definition_for_spec(
    spec: &ToolSpec,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> Option<CodeModeToolDefinition> {
    code_mode_tool_definitions_for_spec(spec, code_mode_input_schema_max_bytes)
        .into_iter()
        .next()
}

fn effective_input_schema_max_bytes(
    mcp_input_schema_max_bytes: Option<usize>,
    code_mode_max_bytes: Option<usize>,
) -> usize {
    codex_code_mode::DEFAULT_INPUT_SCHEMA_MAX_BYTES
        .max(code_mode_max_bytes.unwrap_or_default())
        .max(mcp_input_schema_max_bytes.unwrap_or_default())
}

fn code_mode_tool_definitions_for_spec(
    spec: &ToolSpec,
    code_mode_input_schema_max_bytes: Option<usize>,
) -> Vec<CodeModeToolDefinition> {
    match spec {
        ToolSpec::Function(tool) => {
            let name = tool.name.clone();
            vec![CodeModeToolDefinition {
                tool_name: ToolName::plain(name.clone()),
                name,
                description: tool.description.clone(),
                kind: CodeModeToolKind::Function,
                input_schema: serde_json::to_value(&tool.parameters).ok(),
                input_schema_max_bytes: Some(effective_input_schema_max_bytes(
                    tool.parameters.mcp_input_schema_max_bytes,
                    code_mode_input_schema_max_bytes,
                )),
                output_schema: tool.output_schema.as_ref().map(ToolOutputSchema::to_value),
            }]
        }
        ToolSpec::Freeform(tool) => {
            let name = tool.name.clone();
            vec![CodeModeToolDefinition {
                tool_name: ToolName::plain(name.clone()),
                name,
                description: tool.description.clone(),
                kind: CodeModeToolKind::Freeform,
                input_schema: None,
                input_schema_max_bytes: None,
                output_schema: None,
            }]
        }
        ToolSpec::Namespace(namespace) => namespace
            .tools
            .iter()
            .map(|tool| match tool {
                ResponsesApiNamespaceTool::Function(tool) => {
                    let tool_name = ToolName::namespaced(namespace.name.clone(), tool.name.clone());
                    CodeModeToolDefinition {
                        name: code_mode_name_for_tool_name(&tool_name),
                        tool_name,
                        description: tool.description.clone(),
                        kind: CodeModeToolKind::Function,
                        input_schema: serde_json::to_value(&tool.parameters).ok(),
                        input_schema_max_bytes: Some(effective_input_schema_max_bytes(
                            tool.parameters.mcp_input_schema_max_bytes,
                            code_mode_input_schema_max_bytes,
                        )),
                        output_schema: tool.output_schema.as_ref().map(ToolOutputSchema::to_value),
                    }
                }
                ResponsesApiNamespaceTool::Custom(tool) => {
                    let tool_name = ToolName::namespaced(namespace.name.clone(), tool.name.clone());
                    CodeModeToolDefinition {
                        name: code_mode_name_for_tool_name(&tool_name),
                        tool_name,
                        description: tool.description.clone(),
                        kind: CodeModeToolKind::Freeform,
                        input_schema: None,
                        input_schema_max_bytes: None,
                        output_schema: None,
                    }
                }
            })
            .collect(),
        ToolSpec::ToolSearch { parameters, .. } => {
            let name = spec.name().to_string();
            vec![CodeModeToolDefinition {
                tool_name: ToolName::plain(name.clone()),
                name,
                description: codex_code_mode::TOOL_SEARCH_GUIDANCE.to_string(),
                kind: CodeModeToolKind::Function,
                input_schema: serde_json::to_value(parameters).ok(),
                input_schema_max_bytes: Some(effective_input_schema_max_bytes(
                    parameters.mcp_input_schema_max_bytes,
                    code_mode_input_schema_max_bytes,
                )),
                output_schema: Some(serde_json::json!({
                    "type": "object",
                    "properties": {
                        "tools": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "name": { "type": "string" },
                                    "description": { "type": "string" }
                                },
                                "required": ["name", "description"]
                            }
                        }
                    },
                    "required": ["tools"]
                })),
            }]
        }
        ToolSpec::WebSearch { .. } => Vec::new(),
    }
}

pub fn code_mode_name_for_tool_name(tool_name: &ToolName) -> String {
    if tool_name.is_default_namespace() {
        return tool_name.name.clone();
    }

    match tool_name.namespace.as_deref() {
        Some(namespace) if namespace.ends_with('_') || tool_name.name.starts_with('_') => {
            format!("{namespace}{}", tool_name.name)
        }
        Some(namespace) => format!("{namespace}__{}", tool_name.name),
        None => tool_name.name.clone(),
    }
}

#[cfg(test)]
#[path = "code_mode_tests.rs"]
mod tests;