Support deferred custom tools in tool search (#36998)

## What changed

- Include top-level freeform tools in the tool-search index and mark them for deferred loading.
- Serialize searched freeform tools as Responses API `custom` tools and convert them back into executable tool specs after discovery.
- Cover indexing and the end-to-end search, load, and custom-call routing flow.

GitOrigin-RevId: 124a7705876067542e1bd609546adc9a8d0ef67e

codex-rs/core/tests/suite/search_tool.rs

@@ -5,6 +5,10 @@ use anyhow::Result;
 use codex_config::types::McpServerConfig;
 use codex_config::types::McpServerTransportConfig;
 use codex_core::StartThreadOptions;
+use codex_core::config::Config;
+use codex_extension_api::ExtensionData;
+use codex_extension_api::ExtensionRegistryBuilder;
+use codex_extension_api::ToolContributor;
 use codex_features::Feature;
 use codex_login::CodexAuth;
 use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem;
@@ -20,6 +24,18 @@ use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::McpInvocation;
 use codex_protocol::protocol::Op;
 use codex_protocol::user_input::UserInput;
+use codex_tools::FreeformTool;
+use codex_tools::FreeformToolFormat;
+use codex_tools::FunctionCallError;
+use codex_tools::JsonToolOutput;
+use codex_tools::ToolCall;
+use codex_tools::ToolExecutor;
+use codex_tools::ToolExecutorFuture;
+use codex_tools::ToolExposure;
+use codex_tools::ToolName;
+use codex_tools::ToolOutput;
+use codex_tools::ToolPayload;
+use codex_tools::ToolSpec;
 use core_test_support::apps_test_server::AppsTestServer;
 use core_test_support::apps_test_server::AppsTestToolLoading;
 use core_test_support::apps_test_server::CALENDAR_CREATE_EVENT_MCP_APP_RESOURCE_URI;
@@ -39,6 +55,7 @@ use core_test_support::apps_test_server::search_capable_apps_builder as configur
 use core_test_support::responses::ResponsesRequest;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
+use core_test_support::responses::ev_custom_tool_call;
 use core_test_support::responses::ev_function_call_with_namespace;
 use core_test_support::responses::ev_response_created;
 use core_test_support::responses::ev_tool_search_call;
@@ -56,6 +73,7 @@ use pretty_assertions::assert_eq;
 use serde_json::Value;
 use serde_json::json;
 use std::collections::HashMap;
+use std::sync::Arc;
 use std::time::Duration;
 
 const SEARCH_TOOL_DESCRIPTION_SNIPPETS: [&str; 2] = [
@@ -919,6 +937,126 @@ async fn tool_search_returns_deferred_v1_multi_agent_tools() -> Result<()> {
     Ok(())
 }
 
+struct DeferredCustomTool;
+
+impl ToolContributor for DeferredCustomTool {
+    fn tools(
+        &self,
+        _session_store: &ExtensionData,
+        _thread_store: &ExtensionData,
+    ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
+        vec![Arc::new(Self)]
+    }
+}
+
+impl ToolExecutor<ToolCall> for DeferredCustomTool {
+    fn tool_name(&self) -> ToolName {
+        ToolName::plain("custom_echo")
+    }
+
+    fn spec(&self) -> ToolSpec {
+        ToolSpec::Freeform(FreeformTool {
+            name: "custom_echo".to_string(),
+            description: "Echo a custom payload.".to_string(),
+            defer_loading: None,
+            format: FreeformToolFormat {
+                r#type: "grammar".to_string(),
+                syntax: "lark".to_string(),
+                definition: "start: /.+/".to_string(),
+            },
+        })
+    }
+
+    fn exposure(&self) -> ToolExposure {
+        ToolExposure::Deferred
+    }
+
+    fn handle(&self, call: ToolCall) -> ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            let ToolPayload::Custom { input } = call.payload else {
+                return Err(FunctionCallError::Fatal(
+                    "expected custom tool payload".to_string(),
+                ));
+            };
+            Ok(Box::new(JsonToolOutput::new(json!({
+                "echo": input,
+                "namespace": call.tool_name.namespace,
+            }))) as Box<dyn ToolOutput>)
+        })
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn tool_search_returns_deferred_custom_tool_and_routes_follow_up_call() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let mock = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_tool_search_call("search-1", &json!({ "query": "custom payload" })),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_custom_tool_call("custom-1", "custom_echo", "hello"),
+                ev_completed("resp-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-3"),
+                ev_assistant_message("msg-1", "done"),
+                ev_completed("resp-3"),
+            ]),
+        ],
+    )
+    .await;
+
+    let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+    extensions.tool_contributor(Arc::new(DeferredCustomTool));
+    let mut builder = test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .with_config(configure_search_capable_model);
+    let test = builder.build_with_auto_env(&server).await?;
+    test.submit_turn("Find and run the custom echo tool")
+        .await?;
+
+    let requests = mock.requests();
+    assert_eq!(requests.len(), 3);
+
+    let initial_tools = tool_names(&requests[0].body_json());
+    assert!(
+        initial_tools
+            .iter()
+            .any(|name| name == TOOL_SEARCH_TOOL_NAME)
+    );
+    assert!(initial_tools.iter().all(|name| name != "custom_echo"));
+    assert_eq!(
+        tool_search_output_tools(&requests[1], "search-1"),
+        vec![json!({
+            "type": "custom",
+            "name": "custom_echo",
+            "description": "Echo a custom payload.",
+            "defer_loading": true,
+            "format": {
+                "type": "grammar",
+                "syntax": "lark",
+                "definition": "start: /.+/",
+            },
+        })]
+    );
+    let output = requests[2].custom_tool_call_output("custom-1");
+    let output: Value = serde_json::from_str(
+        output["output"]
+            .as_str()
+            .expect("custom tool output should contain serialized JSON"),
+    )?;
+    assert_eq!(output, json!({ "echo": "hello", "namespace": null }));
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn tool_search_returns_deferred_dynamic_tool_and_routes_follow_up_call() -> Result<()> {
     skip_if_no_network!(Ok(()));

codex-rs/tools/src/responses_api.rs

@@ -46,6 +46,8 @@ pub enum LoadableToolSpec {
     #[allow(dead_code)]
     #[serde(rename = "function")]
     Function(ResponsesApiTool),
+    #[serde(rename = "custom")]
+    Custom(FreeformTool),
     #[serde(rename = "namespace")]
     Namespace(ResponsesApiNamespace),
 }
@@ -85,8 +87,8 @@ pub fn coalesce_loadable_tool_specs(
     let mut coalesced_specs = Vec::new();
     for spec in specs {
         match spec {
-            LoadableToolSpec::Function(tool) => {
+            tool @ (LoadableToolSpec::Function(_) | LoadableToolSpec::Custom(_)) => {
-                coalesced_specs.push(LoadableToolSpec::Function(tool));
+                coalesced_specs.push(tool);
             }
             LoadableToolSpec::Namespace(mut namespace) => {
                 if let Some(existing_namespace) =
@@ -96,7 +98,9 @@ pub fn coalesce_loadable_tool_specs(
                         {
                             Some(existing_namespace)
                         }
-                        LoadableToolSpec::Function(_) | LoadableToolSpec::Namespace(_) => None,
+                        LoadableToolSpec::Function(_)
+                        | LoadableToolSpec::Custom(_)
+                        | LoadableToolSpec::Namespace(_) => None,
                     })
                 {
                     existing_namespace.tools.append(&mut namespace.tools);

codex-rs/tools/src/tool_search.rs

@@ -38,6 +38,10 @@ impl ToolSearchInfo {
                 tool.output_schema = None;
                 LoadableToolSpec::Function(tool)
             }
+            ToolSpec::Freeform(mut tool) => {
+                tool.defer_loading = Some(true);
+                LoadableToolSpec::Custom(tool)
+            }
             ToolSpec::Namespace(mut namespace) => {
                 if namespace.description.trim().is_empty() {
                     namespace.description = default_namespace_description(&namespace.name);
@@ -55,7 +59,7 @@ impl ToolSearchInfo {
                 }
                 LoadableToolSpec::Namespace(namespace)
             }
-            ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } | ToolSpec::Freeform(_) => {
+            ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => {
                 return None;
             }
         };

codex-rs/tools/src/tool_search_tests.rs

@@ -2,6 +2,36 @@ use super::*;
 use pretty_assertions::assert_eq;
 use std::collections::BTreeMap;
 
+#[test]
+fn top_level_custom_tools_are_searchable() {
+    let custom_tool = crate::FreeformTool {
+        name: "apply_patch".to_string(),
+        description: "Apply a patch".to_string(),
+        defer_loading: None,
+        format: crate::FreeformToolFormat {
+            r#type: "grammar".to_string(),
+            syntax: "lark".to_string(),
+            definition: "start: \"patch\"".to_string(),
+        },
+    };
+    let search_info = ToolSearchInfo::from_tool_spec(
+        ToolSpec::Freeform(custom_tool.clone()),
+        /*source_info*/ None,
+    )
+    .expect("top-level custom tool should be searchable");
+
+    assert_eq!(
+        (search_info.entry.search_text, search_info.entry.output),
+        (
+            "apply_patch Apply a patch lark".to_string(),
+            LoadableToolSpec::Custom(crate::FreeformTool {
+                defer_loading: Some(true),
+                ..custom_tool
+            }),
+        )
+    );
+}
+
 #[test]
 fn default_search_text_uses_model_visible_namespace_metadata_once() {
     let mut schedule_schema = JsonSchema::object(

codex-rs/tools/src/tool_spec.rs

@@ -68,6 +68,7 @@ impl From<LoadableToolSpec> for ToolSpec {
     fn from(value: LoadableToolSpec) -> Self {
         match value {
             LoadableToolSpec::Function(tool) => ToolSpec::Function(tool),
+            LoadableToolSpec::Custom(tool) => ToolSpec::Freeform(tool),
             LoadableToolSpec::Namespace(namespace) => ToolSpec::Namespace(namespace),
         }
     }