Keep extension tools directly visible

codex-rs/core/src/tools/handlers/extension_tools.rs

@@ -11,6 +11,7 @@ use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
 use crate::tools::registry::CoreToolRuntime;
 use crate::tools::registry::ToolExecutor;
+use crate::tools::registry::ToolExposure;
 
 pub(crate) struct ExtensionToolAdapter(Arc<dyn codex_tools::ToolExecutor<ExtensionToolCall>>);
 
@@ -27,11 +28,36 @@ impl ToolExecutor<ToolInvocation> for ExtensionToolAdapter {
     }
 
     fn spec(&self) -> ToolSpec {
-        self.0.spec()
+        let mut spec = self.0.spec();
+        if self.0.exposure() == ToolExposure::Deferred {
+            match &mut spec {
+                ToolSpec::Function(tool) => {
+                    tool.defer_loading = None;
+                }
+                ToolSpec::Namespace(namespace) => {
+                    for tool in &mut namespace.tools {
+                        let codex_tools::ResponsesApiNamespaceTool::Function(tool) = tool;
+                        tool.defer_loading = None;
+                    }
+                }
+                ToolSpec::ToolSearch { .. }
+                | ToolSpec::ImageGeneration { .. }
+                | ToolSpec::WebSearch { .. }
+                | ToolSpec::Freeform(_) => {}
+            }
+        }
+        spec
     }
 
-    fn exposure(&self) -> crate::tools::registry::ToolExposure {
-        self.0.exposure()
+    fn exposure(&self) -> ToolExposure {
+        // Extension tools do not yet provide search metadata, so keep them in
+        // the model-visible list even if the shared executor requests deferral.
+        match self.0.exposure() {
+            ToolExposure::Direct => ToolExposure::Direct,
+            ToolExposure::Deferred => ToolExposure::Direct,
+            ToolExposure::DirectModelOnly => ToolExposure::DirectModelOnly,
+            ToolExposure::Hidden => ToolExposure::Hidden,
+        }
     }
 
     fn supports_parallel_tool_calls(&self) -> bool {

codex-rs/core/src/tools/spec_plan_tests.rs

@@ -257,6 +257,49 @@ fn use_bedrock_provider(turn: &mut TurnContext) {
     turn.provider = create_model_provider(provider_info, turn.auth_manager.clone());
 }
 
+struct SpecOnlyExtensionTool {
+    name: &'static str,
+    description: &'static str,
+    exposure: ToolExposure,
+    defer_loading: Option<bool>,
+}
+
+#[async_trait::async_trait]
+impl ToolExecutor<ExtensionToolCall> for SpecOnlyExtensionTool {
+    fn tool_name(&self) -> ToolName {
+        ToolName::plain(self.name)
+    }
+
+    fn spec(&self) -> ToolSpec {
+        ToolSpec::Function(ResponsesApiTool {
+            name: self.name.to_string(),
+            description: self.description.to_string(),
+            strict: true,
+            parameters: codex_tools::JsonSchema::object(
+                BTreeMap::from([(
+                    "message".to_string(),
+                    codex_tools::JsonSchema::string(/*description*/ None),
+                )]),
+                Some(vec!["message".to_string()]),
+                Some(false.into()),
+            ),
+            output_schema: None,
+            defer_loading: self.defer_loading,
+        })
+    }
+
+    fn exposure(&self) -> ToolExposure {
+        self.exposure
+    }
+
+    async fn handle(
+        &self,
+        _call: ExtensionToolCall,
+    ) -> Result<Box<dyn ToolOutput>, codex_tools::FunctionCallError> {
+        panic!("spec planning should not execute extension tools")
+    }
+}
+
 struct WebRunExtensionTool;
 
 #[async_trait::async_trait]
@@ -373,6 +416,65 @@ fn apply_patch_accepts_environment_id(spec: &ToolSpec) -> bool {
     }
 }
 
+#[tokio::test]
+async fn deferred_extension_tools_remain_model_visible() {
+    let plan = probe_with(
+        |turn| {
+            turn.model_info.supports_search_tool = true;
+        },
+        ToolPlanInputs {
+            extension_tool_executors: vec![
+                Arc::new(SpecOnlyExtensionTool {
+                    name: "extension_echo",
+                    description: "Echoes arguments through an extension tool.",
+                    exposure: ToolExposure::Deferred,
+                    defer_loading: None,
+                }),
+                Arc::new(SpecOnlyExtensionTool {
+                    name: "extension_lazy",
+                    description: "Lazy extension tool.",
+                    exposure: ToolExposure::Deferred,
+                    defer_loading: Some(true),
+                }),
+                Arc::new(SpecOnlyExtensionTool {
+                    name: "extension_model_only",
+                    description: "Model-only extension tool.",
+                    exposure: ToolExposure::DirectModelOnly,
+                    defer_loading: None,
+                }),
+            ],
+            ..Default::default()
+        },
+    )
+    .await;
+
+    plan.assert_visible_contains(&["extension_echo", "extension_lazy", "extension_model_only"]);
+    assert_eq!(plan.exposure("extension_echo"), ToolExposure::Direct);
+    assert_eq!(plan.exposure("extension_lazy"), ToolExposure::Direct);
+    assert_eq!(
+        plan.exposure("extension_model_only"),
+        ToolExposure::DirectModelOnly
+    );
+    assert_eq!(
+        plan.visible_spec("extension_lazy"),
+        &ToolSpec::Function(ResponsesApiTool {
+            name: "extension_lazy".to_string(),
+            description: "Lazy extension tool.".to_string(),
+            strict: true,
+            parameters: codex_tools::JsonSchema::object(
+                BTreeMap::from([(
+                    "message".to_string(),
+                    codex_tools::JsonSchema::string(/*description*/ None),
+                )]),
+                Some(vec!["message".to_string()]),
+                Some(false.into()),
+            ),
+            output_schema: None,
+            defer_loading: None,
+        })
+    );
+}
+
 #[tokio::test]
 async fn shell_family_registers_visible_unified_exec_and_hidden_legacy_shell() {
     let plan = probe(|turn| {