From 3947f0d0c3e255bade02e241c16cb43d284c0e65 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 24 Jul 2026 01:17:44 +0000 Subject: [PATCH] Avoid duplicating deferred sources in tool search (#35065) ## Why Deferred tool world state already advertises the available tool sources, so repeating them in the `tool_search` description adds redundant context. ## What changed - Omit the source listing from `tool_search` when `DeferredToolWorldState` is enabled while preserving the tool discovery instructions. - Include the source-listing mode in the tool search handler cache key so the description is rebuilt when the feature state changes. ## Testing Added coverage for source omission, cache invalidation when the feature is toggled, and the resulting request payload. GitOrigin-RevId: f7dc0c4f3351a7c7daadec936b26a080202730bc --- .../core/src/tools/handlers/tool_search.rs | 36 +++++++-- .../src/tools/handlers/tool_search_spec.rs | 80 +++++++++++++------ codex-rs/core/src/tools/spec_plan.rs | 14 +++- codex-rs/core/src/tools/spec_plan_tests.rs | 44 ++++++++++ codex-rs/core/tests/suite/search_tool.rs | 52 ++++++++++++ 5 files changed, 194 insertions(+), 32 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index e5e607a48f..d599376108 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -3,6 +3,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::context::ToolSearchOutput; use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::tool_search_spec::ToolSearchSourceListing; use crate::tools::handlers::tool_search_spec::create_tool_search_tool; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -24,6 +25,7 @@ use tracing::instrument; pub struct ToolSearchHandler { search_infos: Vec, + source_listing: ToolSearchSourceListing, spec: ToolSpec, search_engine: SearchEngine, } @@ -35,20 +37,26 @@ pub(crate) struct ToolSearchHandlerCache { impl ToolSearchHandlerCache { #[instrument(level = "trace", skip_all, fields(search_info_count = search_infos.len()))] - pub(crate) fn get_or_build(&self, search_infos: Vec) -> Arc { + pub(crate) fn get_or_build( + &self, + search_infos: Vec, + source_listing: ToolSearchSourceListing, + ) -> Arc { { let cached = self.cached(); if let Some(cached) = cached.as_ref() && cached.search_infos == search_infos + && cached.source_listing == source_listing { return Arc::clone(cached); } } - let handler = Arc::new(ToolSearchHandler::new(search_infos)); + let handler = Arc::new(ToolSearchHandler::new(search_infos, source_listing)); let mut cached = self.cached(); if let Some(cached) = cached.as_ref() && cached.search_infos == handler.search_infos + && cached.source_listing == handler.source_listing { return Arc::clone(cached); } @@ -71,12 +79,19 @@ impl ToolSearchHandler { skip_all, fields(search_info_count = search_infos.len()) )] - pub(crate) fn new(search_infos: Vec) -> Self { + pub(crate) fn new( + search_infos: Vec, + source_listing: ToolSearchSourceListing, + ) -> Self { let search_source_infos = search_infos .iter() .filter_map(|search_info| search_info.source_info.clone()) .collect::>(); - let spec = create_tool_search_tool(&search_source_infos, TOOL_SEARCH_DEFAULT_LIMIT); + let spec = create_tool_search_tool( + &search_source_infos, + TOOL_SEARCH_DEFAULT_LIMIT, + source_listing, + ); let documents: Vec> = search_infos .iter() .map(|search_info| search_info.entry.search_text.clone()) @@ -88,6 +103,7 @@ impl ToolSearchHandler { Self { search_infos, + source_listing, spec, search_engine, } @@ -205,16 +221,20 @@ mod tests { .expect("MCP handler should return search info"), ]; - let first = cache.get_or_build(search_infos.clone()); - let second = cache.get_or_build(search_infos.clone()); + let first = cache.get_or_build(search_infos.clone(), ToolSearchSourceListing::Include); + let second = cache.get_or_build(search_infos.clone(), ToolSearchSourceListing::Include); assert!(Arc::ptr_eq(&first, &second)); + let without_sources = + cache.get_or_build(search_infos.clone(), ToolSearchSourceListing::Omit); + assert!(!Arc::ptr_eq(&first, &without_sources)); + let mut changed_search_infos = search_infos; changed_search_infos[0] .entry .search_text .push_str(" changed"); - let changed = cache.get_or_build(changed_search_infos); + let changed = cache.get_or_build(changed_search_infos, ToolSearchSourceListing::Omit); assert!(!Arc::ptr_eq(&first, &changed)); } @@ -257,7 +277,7 @@ mod tests { .search_info() .expect("dynamic handler should return search info") })); - let handler = ToolSearchHandler::new(search_infos); + let handler = ToolSearchHandler::new(search_infos, ToolSearchSourceListing::Include); let results = [ &handler.search_infos[0].entry, &handler.search_infos[2].entry, diff --git a/codex-rs/core/src/tools/handlers/tool_search_spec.rs b/codex-rs/core/src/tools/handlers/tool_search_spec.rs index 2d17d131c7..c8026bd305 100644 --- a/codex-rs/core/src/tools/handlers/tool_search_spec.rs +++ b/codex-rs/core/src/tools/handlers/tool_search_spec.rs @@ -4,9 +4,16 @@ use codex_tools::ToolSearchSourceInfo; use codex_tools::ToolSpec; use std::collections::BTreeMap; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ToolSearchSourceListing { + Include, + Omit, +} + pub(crate) fn create_tool_search_tool( searchable_sources: &[ToolSearchSourceInfo], default_limit: usize, + source_listing: ToolSearchSourceListing, ) -> ToolSpec { let properties = BTreeMap::from([ ( @@ -21,33 +28,41 @@ pub(crate) fn create_tool_search_tool( ), ]); - let mut source_descriptions = BTreeMap::new(); - for source in searchable_sources { - source_descriptions - .entry(source.name.clone()) - .and_modify(|existing: &mut Option| { - if existing.is_none() { - *existing = source.description.clone(); - } - }) - .or_insert(source.description.clone()); - } + let source_section = match source_listing { + ToolSearchSourceListing::Include => { + let mut source_descriptions = BTreeMap::new(); + for source in searchable_sources { + source_descriptions + .entry(source.name.clone()) + .and_modify(|existing: &mut Option| { + if existing.is_none() { + *existing = source.description.clone(); + } + }) + .or_insert(source.description.clone()); + } - let source_descriptions = if source_descriptions.is_empty() { - "None currently enabled.".to_string() - } else { - source_descriptions - .into_iter() - .map(|(name, description)| match description { - Some(description) => format!("- {name}: {description}"), - None => format!("- {name}"), - }) - .collect::>() - .join("\n") + let source_descriptions = if source_descriptions.is_empty() { + "None currently enabled.".to_string() + } else { + source_descriptions + .into_iter() + .map(|(name, description)| match description { + Some(description) => format!("- {name}: {description}"), + None => format!("- {name}"), + }) + .collect::>() + .join("\n") + }; + format!( + "\n\nYou have access to tools from the following sources:\n{source_descriptions}\n" + ) + } + ToolSearchSourceListing::Omit => "\n\n".to_string(), }; let description = format!( - "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nYou have access to tools from the following sources:\n{source_descriptions}\nSome of the tools may not have been provided to you upfront, and you should use this tool (`{TOOL_SEARCH_TOOL_NAME}`) to search for the required tools. For MCP tool discovery, always use `{TOOL_SEARCH_TOOL_NAME}` instead of `list_mcp_resources` or `list_mcp_resource_templates`." + "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.{source_section}Some of the tools may not have been provided to you upfront, and you should use this tool (`{TOOL_SEARCH_TOOL_NAME}`) to search for the required tools. For MCP tool discovery, always use `{TOOL_SEARCH_TOOL_NAME}` instead of `list_mcp_resources` or `list_mcp_resource_templates`." ); ToolSpec::ToolSearch { @@ -90,6 +105,7 @@ mod tests { }, ], /*default_limit*/ 8, + ToolSearchSourceListing::Include, ), ToolSpec::ToolSearch { execution: "client".to_string(), @@ -110,4 +126,22 @@ mod tests { } ); } + + #[test] + fn create_tool_search_tool_omits_sources_when_world_state_advertises_them() { + let ToolSpec::ToolSearch { description, .. } = create_tool_search_tool( + &[ToolSearchSourceInfo { + name: "Google Drive".to_string(), + description: Some("Search files and documents.".to_string()), + }], + /*default_limit*/ 8, + ToolSearchSourceListing::Omit, + ) else { + panic!("expected tool search spec"); + }; + + assert!(!description.contains("You have access to tools from the following sources")); + assert!(!description.contains("Google Drive")); + assert!(description.contains("use this tool (`tool_search`) to search")); + } } diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 0b6657d06d..6897f12b60 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -48,6 +48,7 @@ use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHand use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2; use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2; use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2; +use crate::tools::handlers::tool_search_spec::ToolSearchSourceListing; use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_web_search_tool; @@ -950,7 +951,18 @@ fn append_tool_search_executor( return; } - let handler: PlannedRuntime = context.tool_search_handler_cache.get_or_build(search_infos); + let source_listing = if turn_context + .config + .features + .enabled(Feature::DeferredToolWorldState) + { + ToolSearchSourceListing::Omit + } else { + ToolSearchSourceListing::Include + }; + let handler: PlannedRuntime = context + .tool_search_handler_cache + .get_or_build(search_infos, source_listing); planned_tools.add_arc(handler); } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index a0948b65fa..35526c5ab9 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -962,6 +962,50 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { assert!(!second_description.contains("- first: Tools from first.")); } +#[tokio::test] +async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() { + let cache = ToolSearchHandlerCache::default(); + + for world_state_enabled in [false, true, false] { + let (_session, mut turn) = make_session_and_context().await; + turn.model_info.supports_search_tool = true; + set_feature( + &mut turn, + Feature::DeferredToolWorldState, + world_state_enabled, + ); + let turn = Arc::new(turn); + let step_context = StepContext::for_test(Arc::clone(&turn)); + let router = ToolRouter::from_context( + step_context.turn.as_ref(), + &step_context.environments, + step_context.mcp.as_ref(), + ToolRouterParams { + tool_runtimes: vec![mcp_runtime( + "calendar", + "mcp__calendar", + "lookup", + ToolExposure::Deferred, + )], + tool_suggest_candidates: None, + extension_tool_executors: Vec::new(), + dynamic_tools: &[], + }, + &cache, + ); + let plan = ToolPlanProbe::from_router(router); + let ToolSpec::ToolSearch { description, .. } = plan.visible_spec("tool_search") else { + panic!("expected visible tool_search spec"); + }; + + assert_eq!( + description.contains("- calendar: Tools from calendar."), + !world_state_enabled, + "tool search cache should follow the deferred world-state feature" + ); + } +} + #[tokio::test] async fn request_plugin_install_requires_all_discovery_features() { for disabled_feature in [Feature::ToolSuggest, Feature::Apps, Feature::Plugins] { diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 3143ca6f55..69775b72ec 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -5,6 +5,7 @@ use anyhow::Result; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; use codex_core::StartThreadOptions; +use codex_features::Feature; use codex_login::CodexAuth; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; @@ -391,6 +392,57 @@ async fn search_tool_adds_discovery_instructions_to_tool_description() -> Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn search_tool_omits_sources_when_deferred_tool_world_state_is_enabled() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let apps_server = AppsTestServer::mount_searchable(&server).await?; + let mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + + let mut builder = + configured_builder(apps_server.chatgpt_base_url.clone()).with_config(|config| { + config + .features + .enable(Feature::DeferredToolWorldState) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + test.submit_turn_with_approval_and_permission_profile( + "list tools", + AskForApproval::Never, + PermissionProfile::Disabled, + ) + .await?; + + let request = mock.single_request(); + let body = request.body_json(); + let description = tool_search_description(&body).expect("tool_search description should exist"); + assert!( + !SEARCH_TOOL_DESCRIPTION_SNIPPETS + .iter() + .any(|snippet| description.contains(snippet)), + "tool_search should not repeat sources already advertised in world state: {description:?}" + ); + assert!( + request + .message_input_texts("developer") + .into_iter() + .any(|text| text.contains(SEARCH_CALENDAR_NAMESPACE)), + "world state should advertise the deferred Calendar namespace" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn search_tool_hides_apps_tools_without_search() -> Result<()> { skip_if_no_network!(Ok(()));