diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 6fe3c7dfe7..220fb3c11a 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -74,11 +74,11 @@ impl ToolSearchHandlerCache { .entries() .filter(|tool| tool.exposure.is_deferred()) .filter_map(|tool| { - if tool.runtime.immutable_spec().is_some() { - Some(ToolSearchSource::Immutable(Arc::downgrade(&tool.runtime))) + if let Some(runtime) = tool.cached_runtime() { + Some(ToolSearchSource::Immutable(Arc::downgrade(runtime))) } else { - tool.runtime - .search_info() + registry + .search_info(tool) .map(Box::new) .map(ToolSearchSource::Dynamic) } @@ -313,12 +313,12 @@ impl ToolSearchHandler { else { continue; }; - let Some(runtime) = router.tool_runtime(tool_name) else { + let Some(tool) = router.registered_tool(tool_name) else { continue; }; let definitions = prepare_code_mode_tool_definitions( - Some(runtime.as_ref()), - || Cow::Owned(runtime.spec()), + tool.cached_runtime().map(Arc::as_ref), + || Cow::Owned(Arc::unwrap_or_clone(tool.spec())), code_mode_input_schema_max_bytes, tool_description_first, ); @@ -375,10 +375,10 @@ mod tests { #[test] fn cache_reuses_immutable_handlers_and_rebuilds_for_current_registry_changes() { let cache = ToolSearchHandlerCache::default(); - let runtime: Arc = Arc::new( - McpHandler::new(tool_info("calendar", "create_event", "Create events")) - .expect("MCP tool should convert"), - ); + let mut info = tool_info("calendar", "create_event", "Create events"); + info.callable_namespace = "functions".to_string(); + let runtime: Arc = + Arc::new(McpHandler::new(info.clone()).expect("MCP tool should convert")); let mut registry = ToolRegistry::default(); registry.register_trusted_with_exposure(Arc::clone(&runtime), ToolExposure::Deferred); @@ -390,14 +390,46 @@ mod tests { assert!(!Arc::ptr_eq(&first, &without_sources)); let mut replacement_registry = ToolRegistry::default(); - let replacement = Arc::new( - McpHandler::new(tool_info("calendar", "create_event", "Create events")) - .expect("replacement MCP tool should convert"), - ); + let replacement = + Arc::new(McpHandler::new(info).expect("replacement MCP tool should convert")); replacement_registry.register_trusted_with_exposure(replacement, ToolExposure::Deferred); let replacement = cache.get_or_build(&replacement_registry, ToolSearchSourceListing::Omit); assert!(!Arc::ptr_eq(&without_sources, &replacement)); + registry + .apply_functions_namespace_function_prefixes( + Some(&std::collections::BTreeMap::from([( + "create_event".to_string(), + "Guidance.".to_string(), + )])), + [ToolName::plain("create_event")], + ) + .expect("valid function prefixes"); + let prefixed = cache.get_or_build(®istry, ToolSearchSourceListing::Include); + let reused = cache.get_or_build(®istry, ToolSearchSourceListing::Include); + assert!(Arc::ptr_eq(&prefixed, &reused)); + assert_eq!( + prefixed.search_infos[0].entry.search_text, + first.search_infos[0].entry.search_text + ); + let mut expected = first.search("Create events", /*limit*/ 1); + let LoadableToolSpec::Namespace(namespace) = &mut expected[0] else { + panic!("namespace") + }; + let ResponsesApiNamespaceTool::Function(tool) = &mut namespace.tools[0] else { + panic!("function") + }; + tool.description = format!("Guidance.\n\n{}", tool.description); + assert_eq!(prefixed.search("Create events", /*limit*/ 1), expected); + registry + .apply_functions_namespace_function_prefixes(/*prefixes*/ None, []) + .expect("clear prefixes"); + let cleared = cache.get_or_build(®istry, ToolSearchSourceListing::Include); + assert_eq!( + cleared.search("Create events", /*limit*/ 1), + first.search("Create events", /*limit*/ 1) + ); + let mut disabled_registry = ToolRegistry::default(); disabled_registry.register_trusted_with_exposure(runtime, ToolExposure::Direct); let disabled = cache.get_or_build(&disabled_registry, ToolSearchSourceListing::Omit); diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index f51cc4939e..e0b50c48dd 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -139,7 +139,9 @@ impl ToolCallRuntime { .record_tool_call(&call, &source, &step_context); let router = &step_context.tool_router; let supports_parallel = router.tool_supports_parallel(&call); - let tool_runtime = router.tool_runtime(&call.tool_name); + let tool_runtime = router + .registered_tool(&call.tool_name) + .map(|tool| Arc::clone(&tool.runtime)); let finishes_on_cancellation = tool_runtime .as_ref() .is_some_and(|runtime| runtime.finishes_on_cancellation()); diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index da679ce292..ebde5a7bdc 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -38,7 +38,9 @@ use codex_protocol::parse_command::ParsedCommand; use codex_protocol::protocol::EventMsg; use codex_rollout::state_db; use codex_shell_command::parse_command::parse_shell_script; +use codex_tools::FunctionsNamespaceFunctionPrefixes; use codex_tools::ToolName; +use codex_tools::ToolSearchInfo; use codex_tools::ToolSpec; use futures::future::BoxFuture; use indexmap::IndexMap; @@ -303,16 +305,83 @@ pub(crate) struct PostToolUsePayload { pub(crate) struct RegisteredTool { pub(crate) runtime: Arc, pub(crate) exposure: ToolExposure, + /// Per-step description override; never written back to a shared runtime. + pub(crate) model_spec: Option>, +} + +impl RegisteredTool { + /// Direct declarations and Code Mode read this spec; execution uses the runtime. + pub(crate) fn spec(&self) -> Arc { + self.model_spec + .as_ref() + .or_else(|| self.runtime.immutable_spec()) + .cloned() + .unwrap_or_else(|| Arc::new(self.runtime.spec())) + } + + /// Runtime caches describe the original spec, so cannot serve an overridden one. + pub(crate) fn cached_runtime(&self) -> Option<&Arc> { + (self.model_spec.is_none() && self.runtime.immutable_spec().is_some()) + .then_some(&self.runtime) + } } #[derive(Default)] pub struct ToolRegistry { tools: IndexMap, + functions_namespace_function_prefixes: FunctionsNamespaceFunctionPrefixes, first_collision: Option, pub(crate) tool_policy: Arc, } impl ToolRegistry { + pub(crate) fn search_info(&self, tool: &RegisteredTool) -> Option { + let info = tool.runtime.search_info()?; + // A runtime may advertise a different search schema or description. Preserve it, + // along with its ranking text and source, rather than substituting the normal spec. + match self + .functions_namespace_function_prefixes + .prepare(&tool.runtime.tool_name(), || { + info.entry.to_loadable_spec().into() + }) { + Some(spec) => ToolSearchInfo::from_spec(info.entry.search_text, spec, info.source_info), + None => Some(info), + } + } + + /// Apply catalog prefixes to this step's specs, budgeting only reachable tools. + /// The planner supplies the names; this does not change execution or exposure. + pub(crate) fn apply_functions_namespace_function_prefixes( + &mut self, + prefixes: Option<&BTreeMap>, + available_tools: impl IntoIterator, + ) -> Result<(), &'static str> { + let prefixes = FunctionsNamespaceFunctionPrefixes::new(prefixes, available_tools)?; + for tool in self.tools.values_mut() { + tool.model_spec = prefixes + .prepare(&tool.runtime.tool_name(), || tool.runtime.spec()) + .map(Arc::new); + } + self.functions_namespace_function_prefixes = prefixes; + Ok(()) + } + + fn prepare_tool( + &self, + runtime: Arc, + exposure: ToolExposure, + ) -> RegisteredTool { + let model_spec = self + .functions_namespace_function_prefixes + .prepare(&runtime.tool_name(), || runtime.spec()) + .map(Arc::new); + RegisteredTool { + runtime, + exposure, + model_spec, + } + } + pub(crate) fn with_tool_policy(tool_policy: Arc) -> Self { Self { tool_policy, @@ -359,9 +428,10 @@ impl ToolRegistry { if !self.tool_policy.allows(&tool_name) { return; } + let tool = self.prepare_tool(runtime, exposure); match self.tools.entry(tool_name) { Entry::Vacant(entry) => { - entry.insert(RegisteredTool { runtime, exposure }); + entry.insert(tool); } Entry::Occupied(entry) => { let tool_name = entry.key(); @@ -381,8 +451,8 @@ impl ToolRegistry { } let exposure = runtime.exposure(); - self.tools - .shift_insert(0, tool_name, RegisteredTool { runtime, exposure }); + let tool = self.prepare_tool(runtime, exposure); + self.tools.shift_insert(0, tool_name, tool); } pub(crate) fn register_external(&mut self, runtime: Arc) -> bool { @@ -409,9 +479,10 @@ impl ToolRegistry { return false; } + let tool = self.prepare_tool(runtime, exposure); match self.tools.entry(tool_name) { Entry::Vacant(entry) => { - entry.insert(RegisteredTool { runtime, exposure }); + entry.insert(tool); true } Entry::Occupied(entry) => { @@ -503,11 +574,14 @@ impl ToolRegistry { } pub(crate) fn tool(&self, name: &ToolName) -> Option> { - self.tools - .get(&name.clone().with_default_namespace()) + self.registered_tool(name) .map(|tool| Arc::clone(&tool.runtime)) } + pub(crate) fn registered_tool(&self, name: &ToolName) -> Option<&RegisteredTool> { + self.tools.get(&name.clone().with_default_namespace()) + } + #[cfg(test)] pub(crate) fn tool_names_for_test(&self) -> Vec { let mut names = self.tools.keys().cloned().collect::>(); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 8f8635dcfa..a7a192815b 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -11,7 +11,7 @@ use crate::tools::context::ToolPayload; #[cfg(test)] use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::registry::AnyToolResult; -use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::RegisteredTool; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; @@ -249,8 +249,8 @@ impl ToolRouter { .unwrap_or(false) } - pub(crate) fn tool_runtime(&self, tool_name: &ToolName) -> Option> { - self.registry.tool(tool_name) + pub(crate) fn registered_tool(&self, tool_name: &ToolName) -> Option<&RegisteredTool> { + self.registry.registered_tool(tool_name) } #[instrument(level = "trace", skip_all, err)] diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 05a1e2e65a..379209bc7b 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -190,8 +190,8 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow assert_eq!( router - .tool_runtime(&ToolName::plain(parallel_tool_name)) - .map(|runtime| runtime.tool_name()), + .registered_tool(&ToolName::plain(parallel_tool_name)) + .map(|tool| tool.runtime.tool_name()), Some(ToolName::plain(parallel_tool_name)) ); @@ -362,8 +362,8 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { assert!(router.tool_supports_parallel(&call)); assert_eq!( router - .tool_runtime(&call.tool_name) - .map(|runtime| runtime.tool_name()), + .registered_tool(&call.tool_name) + .map(|tool| tool.runtime.tool_name()), Some(call.tool_name.clone()) ); @@ -378,8 +378,8 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { assert!(!router.tool_supports_parallel(&different_server_call)); assert_eq!( router - .tool_runtime(&different_server_call.tool_name) - .map(|runtime| runtime.tool_name()), + .registered_tool(&different_server_call.tool_name) + .map(|tool| tool.runtime.tool_name()), Some(different_server_call.tool_name.clone()) ); @@ -392,7 +392,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { encrypted_function_args: None, }; assert!(!router.tool_supports_parallel(&hidden_call)); - assert!(router.tool_runtime(&hidden_call.tool_name).is_some()); + assert!(router.registered_tool(&hidden_call.tool_name).is_some()); let nested_only_call = ToolCall { tool_name: ToolName::namespaced("mcp__nested_echo__", "query_with_delay"), @@ -524,6 +524,7 @@ fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> RegisteredTool { RegisteredTool { exposure: runtime.exposure(), runtime, + model_spec: None, } } diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 4abd1bbb07..5c83302ca6 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -54,7 +54,6 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_web_search_tool; use crate::tools::multi_agent_tool::multi_agent_v2_handler; -#[cfg(test)] use crate::tools::registry::RegisteredTool; use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; @@ -374,6 +373,7 @@ pub(crate) fn finalize_tool_router( mut hosted_specs: Vec, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> CodexResult { + let model_messages = ResolvedModelMessages::from_model(model_info); hosted_specs.retain(|spec| registry.tool_policy.allows(&ToolName::plain(spec.name()))); apply_direct_model_only_namespace_overrides(turn_context, &mut registry); enforce_strict_3p_tools(turn_context, model_info, &mut registry); @@ -390,13 +390,32 @@ pub(crate) fn finalize_tool_router( } } let tool_search_name = ToolName::plain(TOOL_SEARCH_TOOL_NAME); - if model_info.supports_search_tool + let search_exposure = if model_info.supports_search_tool + && registry.tool_policy.allows(&tool_search_name) && registry.entries().any(|tool| { tool.runtime.tool_name() != tool_search_name && tool.exposure.is_deferred() && tool.runtime.search_info().is_some() - }) - { + }) { + Some( + if turn_context + .config + .features + .enabled(Feature::CodeModeToolSearch) + { + tool_exposure_with_namespace_override( + turn_context, + tool_search_name.clone(), + ToolExposure::Direct, + ) + } else { + ToolExposure::Direct + }, + ) + } else { + None + }; + if search_exposure.is_some() { if registry.remove(&tool_search_name).is_some() { registry.record_collision(tool_search_name.clone()); } @@ -423,10 +442,24 @@ pub(crate) fn finalize_tool_router( registry.record_collision(tool_name); } } - append_tool_search_executor(turn_context, &mut registry, tool_search_handler_cache); } - let model_messages = ResolvedModelMessages::from_model(model_info); + prepare_function_description_prefixes( + turn_context, + model_info, + &mut registry, + search_exposure, + )?; + + if let Some(exposure) = search_exposure { + append_tool_search_executor( + turn_context, + &mut registry, + tool_search_handler_cache, + exposure, + ); + } + let indirect_prefixes = IndirectNamespacePrefixes::new( model_messages.indirect_description_prefixes(), registry.mcp_namespaces(), @@ -621,16 +654,12 @@ fn build_model_visible_specs( let mut specs = Vec::new(); for tool in registry.entries() { let exposure = tool.exposure; - if !exposure.is_direct() { + if !is_directly_visible(turn_context, model_info, tool) { continue; } let tool_name = tool.runtime.tool_name(); - if is_hidden_by_code_mode_only(turn_context, model_info, &tool_name, exposure) { - continue; - } - - let spec = tool.runtime.spec(); + let spec = Arc::unwrap_or_clone(tool.spec()); specs.push(spec_for_model_request( turn_context, model_info, @@ -849,6 +878,20 @@ fn is_hidden_by_code_mode_only( )) } +fn is_directly_visible( + turn_context: &TurnContext, + model_info: &ModelInfo, + tool: &RegisteredTool, +) -> bool { + tool.exposure.is_direct() + && !is_hidden_by_code_mode_only( + turn_context, + model_info, + &tool.runtime.tool_name(), + tool.exposure, + ) +} + fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) -> bool { let tool_name = tool_name.clone().with_default_namespace(); tool_name.namespace.as_ref().is_some_and(|namespace| { @@ -860,6 +903,131 @@ fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) }) } +fn prepare_function_description_prefixes( + turn_context: &TurnContext, + model_info: &ModelInfo, + registry: &mut ToolRegistry, + search_exposure: Option, +) -> CodexResult<()> { + let model_messages = ResolvedModelMessages::from_model(model_info); + let prefixes = model_messages.functions_namespace_functions_description_prefixes(); + let mut available_tools = Vec::new(); + if prefixes.is_some_and(|prefixes| !prefixes.is_empty()) { + let code_mode_enabled = matches!( + effective_tool_mode(turn_context, model_info), + ToolMode::CodeMode | ToolMode::CodeModeOnly + ); + let search_name = ToolName::plain(TOOL_SEARCH_TOOL_NAME); + let search_visible = search_exposure.is_some_and(|exposure| { + !is_hidden_by_code_mode_only(turn_context, model_info, &search_name, exposure) + || (code_mode_enabled + && registry + .tool_policy + .allows(&ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)) + && turn_context + .config + .features + .enabled(Feature::CodeModeToolSearch) + && exposure.is_available_in_code_mode() + && !is_excluded_from_code_mode(turn_context, &search_name)) + }); + available_tools.extend( + registry + .entries() + .filter(|tool| { + is_directly_visible(turn_context, model_info, tool) + || (search_visible + && tool.exposure.is_deferred() + && tool.runtime.search_info().is_some()) + }) + .map(|tool| tool.runtime.tool_name()), + ); + available_tools.extend( + code_mode_tools(turn_context, model_info, registry) + .map(|(_, _, tool)| tool.runtime.tool_name()), + ); + // These wrappers are generated after their inner descriptions are ready. + if code_mode_enabled { + available_tools.extend( + [ + codex_code_mode::PUBLIC_TOOL_NAME, + codex_code_mode::WAIT_TOOL_NAME, + ] + .into_iter() + .map(ToolName::plain) + .filter(|name| registry.tool_policy.allows(name)), + ); + } + } + registry + .apply_functions_namespace_function_prefixes(prefixes, available_tools) + .map_err(|error| CodexErrorDetails::InvalidRequest(error.to_owned()).into()) +} + +/// Select the same nested tools for prefix budgeting and Code Mode construction. +/// A registered tool alone is not reachable when policy disallows the exec wrapper. +fn code_mode_tools<'a>( + turn_context: &'a TurnContext, + model_info: &'a ModelInfo, + registry: &'a ToolRegistry, +) -> impl Iterator, &'a RegisteredTool)> { + let tool_mode = effective_tool_mode(turn_context, model_info); + let enabled = matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) + && registry + .tool_policy + .allows(&ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)); + let mut seen = HashSet::new(); + registry.entries().filter_map(move |tool| { + if !enabled || !tool.exposure.is_available_in_code_mode() { + return None; + } + + let tool_name = tool.runtime.tool_name(); + // Deferred third-party tools still need an exec route to be discoverable. + if is_excluded_from_code_mode(turn_context, &tool_name) + && !(code_mode_only_strict_3p_tools(turn_context, model_info) + && tool.runtime.is_third_party_tool()) + { + return None; + } + + let spec = tool.spec(); + // Derive the name without serializing and augmenting every tool schema. + let code_mode_name = match spec.as_ref() { + ToolSpec::Function(_) | ToolSpec::Freeform(_) => { + codex_tools::code_mode_name_for_tool_name(&tool_name) + } + ToolSpec::Namespace(namespace) if !namespace.tools.is_empty() => { + codex_tools::code_mode_name_for_tool_name(&tool_name) + } + ToolSpec::ToolSearch { .. } + if turn_context + .config + .features + .enabled(Feature::CodeModeToolSearch) => + { + codex_tools::code_mode_name_for_tool_name(&tool_name) + } + ToolSpec::Namespace(_) | ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => { + return None; + } + }; + if !codex_code_mode::is_code_mode_nested_tool(&code_mode_name) { + return None; + } + let code_mode_name = codex_code_mode::normalize_code_mode_identifier(&code_mode_name); + if !seen.insert(code_mode_name.clone()) { + tracing::warn!( + tool_name = %tool_name, + code_mode_name = %code_mode_name, + "skipping tool with a duplicate normalized code-mode name" + ); + return None; + } + Some((code_mode_name, spec, tool)) + }) +} + fn register_code_mode_executors( turn_context: &TurnContext, model_info: &ModelInfo, @@ -877,64 +1045,10 @@ fn register_code_mode_executors( let mut deferred_exec_prompt_tool_specs = Vec::new(); let mut included_deferred_mcp_output_schema = false; let deferred_tools_guidance_enabled = model_info.supports_search_tool; - for tool in registry.entries() { + for (name, spec, tool) in code_mode_tools(turn_context, model_info, registry) { + code_mode_tool_names.insert(name, tool.runtime.tool_name()); let exposure = tool.exposure; - if !exposure.is_available_in_code_mode() { - continue; - } - - let tool_name = tool.runtime.tool_name(); - // Deferred third-party tools still need an exec route to be discoverable. - if is_excluded_from_code_mode(turn_context, &tool_name) - && !(code_mode_only_strict_3p_tools(turn_context, model_info) - && tool.runtime.is_third_party_tool()) - { - continue; - } - - let immutable_spec = tool.runtime.immutable_spec(); - let cached_runtime = immutable_spec.map(|_| Arc::clone(&tool.runtime)); - let spec = immutable_spec - .map(Arc::clone) - .unwrap_or_else(|| Arc::new(tool.runtime.spec())); - // Derive the name without serializing and augmenting every tool schema. - let code_mode_name = match spec.as_ref() { - ToolSpec::Function(_) | ToolSpec::Freeform(_) => { - codex_tools::code_mode_name_for_tool_name(&tool_name) - } - ToolSpec::Namespace(namespace) if !namespace.tools.is_empty() => { - codex_tools::code_mode_name_for_tool_name(&tool_name) - } - ToolSpec::ToolSearch { .. } - if turn_context - .config - .features - .enabled(Feature::CodeModeToolSearch) => - { - codex_tools::code_mode_name_for_tool_name(&tool_name) - } - ToolSpec::Namespace(_) | ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => { - continue; - } - }; - if !codex_code_mode::is_code_mode_nested_tool(&code_mode_name) { - continue; - } - match code_mode_tool_names.entry(codex_code_mode::normalize_code_mode_identifier( - &code_mode_name, - )) { - Entry::Vacant(entry) => { - entry.insert(tool_name); - } - Entry::Occupied(_) => { - tracing::warn!( - tool_name = %tool_name, - code_mode_name = %code_mode_name, - "skipping tool with a duplicate normalized code-mode name" - ); - continue; - } - } + let cached_runtime = tool.cached_runtime().cloned(); if exposure == ToolExposure::Deferred { if deferred_tools_guidance_enabled @@ -1542,6 +1656,7 @@ fn append_tool_search_executor( turn_context: &TurnContext, registry: &mut ToolRegistry, tool_search_handler_cache: &ToolSearchHandlerCache, + exposure: ToolExposure, ) { let source_listing = if turn_context .config @@ -1553,16 +1668,6 @@ fn append_tool_search_executor( ToolSearchSourceListing::Include }; let handler = tool_search_handler_cache.get_or_build(registry, source_listing); - // Preserve native search's exposure when nested search is disabled. - let exposure = if turn_context - .config - .features - .enabled(Feature::CodeModeToolSearch) - { - tool_exposure_with_namespace_override(turn_context, handler.tool_name(), handler.exposure()) - } else { - handler.exposure() - }; registry.register_trusted_with_exposure(handler, exposure); } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 395cc95e5d..1ccb81393b 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -430,6 +430,7 @@ fn mcp_runtime( RegisteredTool { runtime: handler, exposure, + model_spec: None, } } @@ -1766,6 +1767,7 @@ async fn strict_namespace_ownership_requires_tool_namespace_inventory_opt_in() { RegisteredTool { runtime: Arc::new(McpHandler::new(tool).expect("MCP tool spec should build")), exposure, + model_spec: None, } }) .collect(); @@ -2062,6 +2064,7 @@ async fn strict_tool_collisions_allow_multiple_tools_in_one_namespace() { McpHandler::new(undocumented_tool).expect("MCP tool spec should build"), ), exposure: ToolExposure::Direct, + model_spec: None, }, mcp_runtime("shared", "shared", "lookup", ToolExposure::Direct), mcp_runtime("shared", "shared", "list", ToolExposure::Direct), @@ -2101,6 +2104,7 @@ async fn relaxed_tool_collisions_preserve_first_nonempty_namespace_description() RegisteredTool { runtime: Arc::new(McpHandler::new(tool).expect("MCP tool spec should build")), exposure: ToolExposure::Direct, + model_spec: None, } }; let plan = probe_with( diff --git a/codex-rs/core/src/unified_exec/shell_snapshot.rs b/codex-rs/core/src/unified_exec/shell_snapshot.rs index dc59f2c999..c7cd928984 100644 --- a/codex-rs/core/src/unified_exec/shell_snapshot.rs +++ b/codex-rs/core/src/unified_exec/shell_snapshot.rs @@ -38,7 +38,7 @@ impl Session { if !self.features().enabled(Feature::ShellSnapshotV2) || step_context .tool_router - .tool_runtime(&ToolName::plain("exec_command")) + .registered_tool(&ToolName::plain("exec_command")) .is_none() { return None; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index fc4b9a9b60..40cd9bd5e5 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -8365,8 +8365,13 @@ async fn code_mode_renders_local_refs_in_outbound_exec_description() -> Result<( Ok(()) } +#[test_case("codex_app", "codex_app__hidden_dynamic_tool"; "named_namespace")] +#[test_case("functions", "hidden_dynamic_tool"; "functions_namespace")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_search_normalizes_and_calls_deferred_dynamic_tools() -> Result<()> { +async fn code_mode_search_normalizes_and_calls_deferred_dynamic_tools( + namespace: &str, + code_mode_name: &str, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -8374,6 +8379,16 @@ async fn code_mode_search_normalizes_and_calls_deferred_dynamic_tools() -> Resul .with_model_info_override("gpt-5.5", |model_info| { model_info.supports_search_tool = true; model_info.tool_mode = Some(ToolMode::CodeMode); + model_info + .model_messages + .get_or_insert_default() + .tools + .get_or_insert_default() + .functions_namespace_functions_description_prefixes = + Some(std::collections::BTreeMap::from([( + "hidden-dynamic-tool".to_string(), + "Function guidance.".to_string(), + )])); }) .with_config(|config| { let _ = config.features.enable(Feature::CodeMode); @@ -8398,7 +8413,7 @@ async fn code_mode_search_normalizes_and_calls_deferred_dynamic_tools() -> Resul )], }), DynamicToolSpec::Namespace(DynamicToolNamespaceSpec { - name: "codex_app".to_string(), + name: namespace.to_string(), description: "Codex app tools.".to_string(), tools: vec![DynamicToolNamespaceTool::Function( DynamicToolFunctionSpec { @@ -8445,6 +8460,17 @@ text( "#; responses::mount_sse_once( + &server, + sse(vec![ + responses::ev_tool_search_call( + "find", + &serde_json::json!({"query": "hidden-dynamic-tool"}), + ), + ev_completed("search"), + ]), + ) + .await; + let search_response = responses::mount_sse_once( &server, sse(vec![ ev_response_created("resp-1"), @@ -8504,7 +8530,7 @@ text( _ => None, }) .await; - assert_eq!(request.namespace.as_deref(), Some("codex_app")); + assert_eq!(request.namespace.as_deref(), Some(namespace)); assert_eq!(request.tool, "hidden-dynamic-tool"); assert_eq!(request.arguments, serde_json::json!({ "city": "Paris" })); test.codex @@ -8538,13 +8564,35 @@ text( )?; assert_eq!( parsed.get("name"), - Some(&Value::String("codex_app__hidden_dynamic_tool".to_string())) + Some(&Value::String(code_mode_name.to_string())) ); assert_eq!( parsed.get("out"), Some(&Value::String("hidden-ok".to_string())) ); assert_eq!(parsed["recoveredFromInvalidArguments"], true); + let loaded = search_response.single_request().tool_search_output("find"); + let loaded = loaded["tools"] + .as_array() + .expect("loaded tools") + .iter() + .find(|tool| tool["name"] == namespace) + .expect("loaded namespace"); + assert_eq!( + loaded["tools"][0]["description"], + if namespace == "functions" { + "Function guidance.\n\nA hidden dynamic tool." + } else { + "A hidden dynamic tool." + } + ); + assert_eq!( + parsed["description"] + .as_str() + .expect("tool description") + .contains("Function guidance.\n\nA hidden dynamic tool."), + namespace == "functions", + ); assert!( parsed .get("description") @@ -8553,7 +8601,7 @@ text( description.contains("Codex app tools.") && description.contains("A hidden dynamic tool.") && description.contains("declare const tools:") - && description.contains("codex_app__hidden_dynamic_tool(args:") + && description.contains(&format!("{code_mode_name}(args:")) }) ); diff --git a/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs b/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs index aeb1b0c197..00a2ddd15b 100644 --- a/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs +++ b/codex-rs/core/tests/suite/scenarios_indirect_namespace_prefixes.rs @@ -1,4 +1,4 @@ -//! Catalog namespace and MCP server prefixes follow the current model without changing dispatch. +//! Catalog function and namespace prefixes follow the current model without changing dispatch. use super::super::code_mode::custom_tool_output_last_non_empty_text; use super::super::rmcp_client::remote_aware_environment_id; @@ -89,6 +89,10 @@ async fn catalog_namespace_prefixes_follow_the_selected_model( model.supports_search_tool = include_search; model.model_messages.as_mut().expect("model messages").tools = Some( serde_json::from_value(json!({ + "functions_namespace_functions_description_prefixes": if slug == "gpt-5.5" { + json!({"exec": "Exec guidance.", "view_image": "Image guidance.", + "list_agents": "Wrong namespace.", "echo": "Wrong namespace."}) + } else { json!({}) }, "indirect_description_prefixes": { "namespaces": prefixes, "mcp_servers": {"reports.one": mcp_prefix}, @@ -169,6 +173,7 @@ text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0] .expect("exec"), ); let description = exec["description"].as_str().expect("description"); + assert_eq!(description.starts_with("Exec guidance.\n\n"), index == 0); for &prefix in &prefixes { assert_eq!( description.matches(prefix).count(), @@ -181,6 +186,19 @@ text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0] .iter() .all(|prefix| !json!(tools).to_string().contains(*prefix)) ); + if index == 0 { + for tool in &mut tools { + if tool["name"] == "view_image" { + tool["description"] = json!( + tool["description"] + .as_str() + .expect("tool description") + .strip_prefix("Image guidance.\n\n") + .expect("function prefix") + ); + } + } + } tools }); assert_eq!(direct_specs[0], direct_specs[1]); @@ -207,15 +225,18 @@ text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0] outputs.push(output); } let mut expected = outputs[1].clone(); - for (description, prefix) in expected["descriptions"] - .as_array_mut() - .expect("tool descriptions") - .iter_mut() - .zip(prefixes) - { - *description = json!(format!( + expected["descriptions"][0] = json!(format!( + "Image guidance.\n\n{}", + expected["descriptions"][0] + .as_str() + .expect("tool description") + )); + for (index, prefix) in prefixes.iter().enumerate() { + expected["descriptions"][index] = json!(format!( "{prefix}\n\n{}", - description.as_str().expect("tool description") + expected["descriptions"][index] + .as_str() + .expect("tool description") )); } assert_eq!(outputs[0], expected); @@ -223,7 +244,7 @@ text({descriptions, echo: report.structuredContent.echo, agent: agents.agents[0] insta::assert_snapshot!( "catalog_indirect_namespace_prefixes", context_snapshot::format_request_history_snapshot( - "Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch.", + "Function prefixes annotate direct and indirect tools; namespace prefixes annotate embedded namespace docs, loaded namespaces, and flat indirect entries. A model switch clears both without changing dispatch.", &requests, &ContextSnapshotOptions::default() .rewrite_known_segments() diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap index c7761cf293..5a09b5f5e3 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__indirect_namespace_prefixes__catalog_indirect_namespace_prefixes.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/scenarios_indirect_namespace_prefixes.rs -expression: "context_snapshot::format_request_history_snapshot(\"Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +expression: "context_snapshot::format_request_history_snapshot(\"Function prefixes annotate direct and indirect tools; namespace prefixes annotate embedded namespace docs, loaded namespaces, and flat indirect entries. A model switch clears both without changing dispatch.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" --- -Scenario: Catalog namespace and raw MCP server prefixes annotate embedded docs and ALL_TOOLS; a model switch clears them without changing nested dispatch. +Scenario: Function prefixes annotate direct and indirect tools; namespace prefixes annotate embedded namespace docs, loaded namespaces, and flat indirect entries. A model switch clears both without changing dispatch. ## Window 1 Settings: @@ -15,8 +15,8 @@ Settings: stream: true text: {"verbosity":"low"} tool_choice: "auto" - tools (4; hash=7707adf4d5cde816): - - custom/exec: Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided JavaScript code in...; hash=66e0c24dd8b9010d + tools (4; hash=1261f9f79ad97cee): + - custom/exec: Exec guidance. Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided Ja...; hash=db96605f80dbbdd7 - function/wait: Waits on a yielded `exec` cell and returns new output or completion. - Use `wait` only after `exe...; args=[cell_id, max_tokens, terminate, yield_time_ms]; hash=6692d8f6c55be489 - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 - web_search; hash=e214d4db1bc49e89 @@ -44,7 +44,7 @@ Settings: 07:custom_tool_call_output:Script completed Wall time seconds Output: - | {"descriptions":["Functions context.\\n\\nView a local image file from the filesystem when visual inspecti...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=22aad8f340003b22] + | {"descriptions":["Functions context.\\n\\nImage guidance.\\n\\nView a local image file from the filesystem w...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=34de0ce5c435f70f] ## Window 2 (after request 2) Settings: relative to window 1 @@ -74,7 +74,7 @@ Settings: relative to window 1 07:custom_tool_call_output:Script completed Wall time seconds Output: - | {"descriptions":["Functions context.\\n\\nView a local image file from the filesystem when visual inspecti...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=22aad8f340003b22] + | {"descriptions":["Functions context.\\n\\nImage guidance.\\n\\nView a local image file from the filesystem w...>>; };\\n```"],"echo":"ECHOING: ping","agent":"/root"} [hash=34de0ce5c435f70f] 08:message/assistant: Done. 09:message/developer: diff --git a/codex-rs/prompts/src/model_messages.rs b/codex-rs/prompts/src/model_messages.rs index 8b043aa3e6..b7d819e776 100644 --- a/codex-rs/prompts/src/model_messages.rs +++ b/codex-rs/prompts/src/model_messages.rs @@ -204,6 +204,17 @@ impl<'a> ResolvedModelMessages<'a> { .as_ref() } + /// Selects per-function guidance restricted to the default `functions` namespace. + pub fn functions_namespace_functions_description_prefixes( + &self, + ) -> Option<&'a std::collections::BTreeMap> { + self.catalog_messages? + .tools + .as_ref()? + .functions_namespace_functions_description_prefixes + .as_ref() + } + /// Selects Code Mode messages; bundled text and runtime composition belong to the tool owner. pub fn code_mode(&self) -> Option<&'a CodeModeToolMessages> { self.catalog_messages?.tools.as_ref()?.code_mode.as_ref() diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 45654aa8c5..1a203e1972 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -580,9 +580,18 @@ pub struct ConfirmationPolicies { pub computer_use: Option, } -/// Model-owned tool messages and indirect namespace guidance. +/// Model-owned tool messages and description prefixes. #[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] pub struct ToolMessages { + /// Prefixes for individual functions in the `functions` namespace, keyed by unqualified name. + /// Applies to direct declarations, Code Mode, and tool search, including custom tools. + /// Values are trimmed; empty values and unavailable functions add nothing. The trimmed + /// values for available functions plus separating blank lines must fit in 256 UTF-8 bytes + /// in total. This bounds the added text, not the complete definition or its token count. + /// Other namespaces are never affected. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub functions_namespace_functions_description_prefixes: + Option>, /// Optional guidance for indirectly presented tools; missing or empty adds no prefix. #[serde(default, skip_serializing_if = "Option::is_none")] pub indirect_description_prefixes: Option, @@ -597,6 +606,7 @@ pub struct ToolMessages { } /// Plain-text prefixes for Code Mode documentation, ALL_TOOLS, and loaded tool-search namespaces. +/// Flat indirect entries carry these before any per-function prefix. /// Values are trimmed, and selectors for the same final namespace must agree, including empty values. /// Empty values add no prefix; unregistered targets are ignored. #[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] diff --git a/codex-rs/tools/src/functions_namespace_function_prefixes.rs b/codex-rs/tools/src/functions_namespace_function_prefixes.rs new file mode 100644 index 0000000000..41755c8091 --- /dev/null +++ b/codex-rs/tools/src/functions_namespace_function_prefixes.rs @@ -0,0 +1,83 @@ +//! Prepares function descriptions before direct, Code Mode, or search rendering. +//! Only the `functions` namespace is affected; runtime-owned specs stay unchanged. + +use crate::ResponsesApiNamespaceTool; +use crate::ToolName; +use crate::ToolSpec; +use std::collections::BTreeMap; + +/// Trimmed, size-bounded description prefixes for the available `functions.*` tools. +#[derive(Default)] +pub struct FunctionsNamespaceFunctionPrefixes(BTreeMap); + +impl FunctionsNamespaceFunctionPrefixes { + pub fn new( + prefixes: Option<&BTreeMap>, + available_tools: impl IntoIterator, + ) -> Result { + let Some(prefixes) = prefixes.filter(|prefixes| !prefixes.is_empty()) else { + return Ok(Self::default()); + }; + let mut remaining = 256usize; + let mut normalized = BTreeMap::new(); + for tool in available_tools { + if !tool.is_default_namespace() || normalized.contains_key(&tool.name) { + continue; + } + let Some(prefix) = prefixes.get(&tool.name) else { + continue; + }; + let prefix = prefix.trim(); + if prefix.is_empty() { + continue; + } + remaining = remaining + .checked_sub(prefix.len()) + .and_then(|remaining| remaining.checked_sub(2)) + .ok_or( + "Available functions' description prefixes exceed the 256-byte combined limit", + )?; + normalized.insert(tool.name, prefix.to_owned()); + } + Ok(Self(normalized)) + } + + /// Returns an owned replacement only for a configured function. Unaffected tools + /// keep their original specs and derived caches, including large MCP schemas. + pub fn prepare(&self, name: &ToolName, spec: impl FnOnce() -> ToolSpec) -> Option { + if !name.is_default_namespace() || !self.0.contains_key(&name.name) { + return None; + } + let mut spec = spec(); + match &mut spec { + ToolSpec::Function(tool) => self.prepend(&tool.name, &mut tool.description), + ToolSpec::Freeform(tool) => self.prepend(&tool.name, &mut tool.description), + ToolSpec::Namespace(namespace) if namespace.name == "functions" => { + for tool in &mut namespace.tools { + match tool { + ResponsesApiNamespaceTool::Function(tool) => { + self.prepend(&tool.name, &mut tool.description); + } + ResponsesApiNamespaceTool::Custom(tool) => { + self.prepend(&tool.name, &mut tool.description); + } + } + } + } + ToolSpec::Namespace(_) | ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => { + return None; + } + } + Some(spec) + } + + fn prepend(&self, name: &str, description: &mut String) { + if let Some(prefix) = self.0.get(name) { + *description = if description.is_empty() { + prefix.clone() + } else { + format!("{prefix}\n\n{description}") + }; + } + } +} diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index b07c7b0cd2..29c02afabb 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -4,6 +4,7 @@ mod code_mode; mod dynamic_tool; mod function_call_error; +mod functions_namespace_function_prefixes; mod image_detail; mod indirect_namespace_prefixes; mod json_schema; @@ -30,6 +31,7 @@ pub use code_mode::tool_spec_to_code_mode_tool_definition; pub use codex_protocol::ToolName; pub use dynamic_tool::parse_dynamic_tool; pub use function_call_error::FunctionCallError; +pub use functions_namespace_function_prefixes::FunctionsNamespaceFunctionPrefixes; pub use image_detail::can_request_original_image_detail; pub use image_detail::normalize_output_image_detail; pub use image_detail::sanitize_original_image_detail;