codex/codex-rs/core/src/unavailable_tool.rs

use std::collections::BTreeMap;
use std::collections::HashSet;

use crate::tools::flat_tool_name;
use codex_protocol::models::ResponseItem;
use codex_tools::ToolName;

pub(crate) fn collect_unavailable_called_tools(
    input: &[ResponseItem],
    exposed_tool_names: &HashSet<ToolName>,
) -> Vec<ToolName> {
    let mut unavailable_tools = BTreeMap::new();
    let exposed_flat_names = exposed_tool_names
        .iter()
        .map(flat_tool_name)
        .collect::<HashSet<_>>();

    for item in input {
        let ResponseItem::FunctionCall {
            name, namespace, ..
        } = item
        else {
            continue;
        };
        if !should_collect_unavailable_tool(name, namespace.as_deref()) {
            continue;
        }

        let tool_name = match namespace {
            Some(namespace) => ToolName::namespaced(namespace.clone(), name.clone()),
            None => ToolName::plain(name.clone()),
        };
        let tool_name_flat = flat_tool_name(&tool_name).into_owned();
        if exposed_flat_names.contains(tool_name_flat.as_str()) {
            continue;
        }

        unavailable_tools
            .entry(tool_name_flat)
            .or_insert_with(|| tool_name);
    }

    unavailable_tools.into_values().collect()
}

fn should_collect_unavailable_tool(name: &str, namespace: Option<&str>) -> bool {
    namespace.is_some_and(|namespace| namespace.starts_with("mcp__")) || name.starts_with("mcp__")
}

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

    fn function_call(name: &str, namespace: Option<&str>) -> ResponseItem {
        ResponseItem::FunctionCall {
            id: None,
            name: name.to_string(),
            namespace: namespace.map(str::to_string),
            arguments: "{}".to_string(),
            call_id: format!("call-{name}"),
        }
    }

    #[test]
    fn collect_unavailable_called_tools_detects_mcp_function_calls() {
        let input = vec![
            function_call("shell", /*namespace*/ None),
            function_call("mcp__server__lookup", /*namespace*/ None),
            function_call("_create_event", Some("mcp__codex_apps__calendar")),
        ];

        let tools = collect_unavailable_called_tools(&input, &HashSet::new());

        assert_eq!(
            tools,
            vec![
                ToolName::namespaced("mcp__codex_apps__calendar", "_create_event"),
                ToolName::plain("mcp__server__lookup"),
            ]
        );
    }

    #[test]
    fn collect_unavailable_called_tools_skips_currently_available_tools() {
        let exposed_tool_names = HashSet::from([
            ToolName::plain("mcp__server__lookup"),
            ToolName::plain("mcp__server__search"),
        ]);
        let input = vec![
            function_call("mcp__server__lookup", /*namespace*/ None),
            function_call("mcp__server__search", /*namespace*/ None),
            function_call("mcp__server__missing", /*namespace*/ None),
        ];

        let tools = collect_unavailable_called_tools(&input, &exposed_tool_names);

        assert_eq!(tools, vec![ToolName::plain("mcp__server__missing")]);
    }

    #[test]
    fn collect_unavailable_called_tools_matches_exposed_display_names() {
        let exposed_tool_names = HashSet::from([ToolName::namespaced("mcp__server__", "lookup")]);
        let input = vec![function_call(
            "mcp__server__lookup",
            /*namespace*/ None,
        )];

        let tools = collect_unavailable_called_tools(&input, &exposed_tool_names);

        assert_eq!(tools, Vec::new());
    }
}