From a6b99ee5c45e35165c9bac9da4cec65f2ca55bc9 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Thu, 9 Jul 2026 08:17:13 -0700 Subject: [PATCH] code-mode: retain shared MCP types for deferred tools (#31745) ## Why When MCP tools are deferred behind `tool_search`, Code mode keeps them callable but omits their individual declarations from the initial `exec` description. The shared MCP `CallToolResult` types were derived only from directly rendered tools, so deferring every MCP tool also removed the common response contract that models need to interpret MCP results. This restores that contract without undoing the context savings from deferred tool definitions. This is a follow-up to #29486. ## What changed - Track deferred Code-mode tool definitions separately from directly rendered definitions. - Render the shared MCP type preamble when either direct or deferred MCP tools are available. - Keep deferred tool declarations out of the initial prompt. - Add unit and integration coverage for deferred MCP tools. ## Testing - `just test -p codex-code-mode-protocol` - `just test -p codex-core code_mode_only_guides_all_tools_search_and_calls_deferred_app_tools` --- .../code-mode-protocol/src/description.rs | 78 ++++++++++++++----- .../core/src/tools/code_mode/execute_spec.rs | 8 +- codex-rs/core/src/tools/spec_plan.rs | 13 ++-- codex-rs/core/tests/suite/code_mode.rs | 1 + 4 files changed, 69 insertions(+), 31 deletions(-) diff --git a/codex-rs/code-mode-protocol/src/description.rs b/codex-rs/code-mode-protocol/src/description.rs index 6ae31ff209..d0504f54a4 100644 --- a/codex-rs/code-mode-protocol/src/description.rs +++ b/codex-rs/code-mode-protocol/src/description.rs @@ -250,25 +250,32 @@ pub fn is_code_mode_nested_tool(tool_name: &str) -> bool { pub fn build_exec_tool_description( enabled_tools: &[ToolDefinition], + deferred_tools: &[ToolDefinition], namespace_descriptions: &BTreeMap, code_mode_only: bool, - deferred_tools_available: bool, ) -> String { let mut sections = Vec::new(); sections.push(EXEC_DESCRIPTION_TEMPLATE.to_string()); - if deferred_tools_available { + if !deferred_tools.is_empty() { sections.push(DEFERRED_NESTED_TOOLS_GUIDANCE.to_string()); } if !code_mode_only { return sections.join("\n\n"); } + let has_mcp_tools = enabled_tools + .iter() + .chain(deferred_tools) + .any(|tool| mcp_structured_content_schema(tool.output_schema.as_ref()).is_some()); + if has_mcp_tools { + sections.push(format!( + "Shared MCP Types:\n```ts\n{MCP_TYPESCRIPT_PREAMBLE}\n```" + )); + } + if !enabled_tools.is_empty() { let mut current_namespace: Option<&str> = None; let mut nested_tool_sections = Vec::with_capacity(enabled_tools.len()); - let has_mcp_tools = enabled_tools - .iter() - .any(|tool| mcp_structured_content_schema(tool.output_schema.as_ref()).is_some()); for tool in enabled_tools { let name = tool.name.as_str(); @@ -305,11 +312,6 @@ pub fn build_exec_tool_description( } } - if has_mcp_tools { - sections.push(format!( - "Shared MCP Types:\n```ts\n{MCP_TYPESCRIPT_PREAMBLE}\n```" - )); - } let nested_tool_reference = nested_tool_sections.join("\n\n"); sections.push(nested_tool_reference); } @@ -863,9 +865,9 @@ mod tests { input_schema: None, output_schema: None, }], + &[], &BTreeMap::new(), /*code_mode_only*/ true, - /*deferred_tools_available*/ false, ); assert!(description.contains( "### `foo` @@ -876,12 +878,8 @@ bar" #[test] fn exec_description_mentions_timeout_helpers() { - let description = build_exec_tool_description( - &[], - &BTreeMap::new(), - /*code_mode_only*/ false, - /*deferred_tools_available*/ false, - ); + let description = + build_exec_tool_description(&[], &[], &BTreeMap::new(), /*code_mode_only*/ false); assert!(description.contains("`setTimeout(callback: () => void, delayMs?: number)`")); assert!(description.contains("`clearTimeout(timeoutId?: number)`")); } @@ -930,9 +928,9 @@ bar" }))), }, ], + &[], &namespace_descriptions, /*code_mode_only*/ true, - /*deferred_tools_available*/ false, ); assert_eq!(description.matches("## mcp__sample").count(), 1); assert!(description.contains("## mcp__sample\nShared namespace guidance.")); @@ -970,9 +968,9 @@ bar" "additionalProperties": false }))), }], + &[], &namespace_descriptions, /*code_mode_only*/ true, - /*deferred_tools_available*/ false, ); assert!(!description.contains("## mcp__sample")); @@ -1069,9 +1067,9 @@ bar" output_schema: second_tool.output_schema, }, ], + &[], &BTreeMap::new(), /*code_mode_only*/ true, - /*deferred_tools_available*/ false, ); assert_eq!( @@ -1083,13 +1081,51 @@ bar" assert_eq!(description.matches("Shared MCP Types:").count(), 1); } + #[test] + fn code_mode_only_description_renders_shared_mcp_types_for_deferred_tools() { + let deferred_tool = ToolDefinition { + name: "mcp__sample__alpha".to_string(), + tool_name: ToolName::namespaced("mcp__sample__", "alpha"), + description: "Deferred tool".to_string(), + kind: CodeModeToolKind::Function, + input_schema: Some(json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + })), + output_schema: Some(mcp_call_tool_result_schema(json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + }))), + }; + + let description = build_exec_tool_description( + &[], + &[deferred_tool], + &BTreeMap::new(), + /*code_mode_only*/ true, + ); + + assert!(description.contains("Some deferred nested tools may be omitted")); + assert!(description.contains("Shared MCP Types:")); + assert!(!description.contains("### `mcp__sample__alpha`")); + } + #[test] fn exec_description_mentions_deferred_nested_tools_when_available() { let description = build_exec_tool_description( &[], + &[ToolDefinition { + name: "deferred_tool".to_string(), + tool_name: ToolName::plain("deferred_tool"), + description: "Deferred tool".to_string(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }], &BTreeMap::new(), /*code_mode_only*/ false, - /*deferred_tools_available*/ true, ); assert!(description.contains("Some deferred nested tools may be omitted")); diff --git a/codex-rs/core/src/tools/code_mode/execute_spec.rs b/codex-rs/core/src/tools/code_mode/execute_spec.rs index 0a858bd206..39fbf69eb4 100644 --- a/codex-rs/core/src/tools/code_mode/execute_spec.rs +++ b/codex-rs/core/src/tools/code_mode/execute_spec.rs @@ -6,9 +6,9 @@ use std::collections::BTreeMap; pub(crate) fn create_code_mode_tool( enabled_tools: &[CodeModeToolDefinition], + deferred_tools: &[CodeModeToolDefinition], namespace_descriptions: &BTreeMap, code_mode_only: bool, - deferred_tools_available: bool, ) -> ToolSpec { const CODE_MODE_FREEFORM_GRAMMAR: &str = r#" start: pragma_source | plain_source @@ -24,9 +24,9 @@ SOURCE: /[\s\S]+/ name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), description: codex_code_mode::build_exec_tool_description( enabled_tools, + deferred_tools, namespace_descriptions, code_mode_only, - deferred_tools_available, ), format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -56,17 +56,17 @@ mod tests { assert_eq!( create_code_mode_tool( &enabled_tools, + &[], &BTreeMap::new(), /*code_mode_only*/ true, - /*deferred_tools_available*/ false, ), ToolSpec::Freeform(FreeformTool { name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), description: codex_code_mode::build_exec_tool_description( &enabled_tools, + &[], &BTreeMap::new(), /*code_mode_only*/ true, - /*deferred_tools_available*/ false ), format: FreeformToolFormat { r#type: "grammar".to_string(), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 57879aeded..f6e50984ac 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -461,7 +461,7 @@ fn build_code_mode_executors( let mut code_mode_nested_tool_specs = Vec::new(); let mut exec_prompt_tool_specs = Vec::new(); - let mut deferred_tools_available = false; + let mut deferred_exec_prompt_tool_specs = Vec::new(); let deferred_tools_guidance_enabled = search_tool_enabled(turn_context); for executor in executors { let exposure = executor.exposure(); @@ -480,10 +480,9 @@ fn build_code_mode_executors( let spec = executor.spec(); if exposure == ToolExposure::Deferred { - // Only show deferred-tool guidance when supported and an included spec is usable by code mode. - deferred_tools_available |= deferred_tools_guidance_enabled - && !collect_code_mode_exec_prompt_tool_definitions(std::iter::once(&spec)) - .is_empty(); + if deferred_tools_guidance_enabled { + deferred_exec_prompt_tool_specs.push(spec.clone()); + } } else { exec_prompt_tool_specs.push(spec.clone()); } @@ -495,14 +494,16 @@ fn build_code_mode_executors( collect_code_mode_exec_prompt_tool_definitions(exec_prompt_tool_specs.iter()); enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); + let deferred_tools = + collect_code_mode_exec_prompt_tool_definitions(deferred_exec_prompt_tool_specs.iter()); vec![ Arc::new(CodeModeExecuteHandler::new( create_code_mode_tool( &enabled_tools, + &deferred_tools, &namespace_descriptions, tool_mode == ToolMode::CodeModeOnly, - deferred_tools_available, ), code_mode_nested_tool_specs, )), diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 3e9d769310..52af3824d5 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -743,6 +743,7 @@ if (!tool) { }) .expect("exec description should be present"); assert!(exec_description.contains("filter `ALL_TOOLS` by `name` and `description`")); + assert!(exec_description.contains("Shared MCP Types:")); assert!(!exec_description.contains("calendar_timezone_option_99")); let request = follow_up_mock.single_request();