openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

register MCP tools lazily through tool search

Browse Source
This commit is contained in:
Ahmed Ibrahim
2026-06-09 16:46:08 +00:00
parent 3c66a4a820
commit e4e781af09
22 changed files with 695 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp_tool_call.rs
+1
View File
@@ -681,6 +681,7 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &
match mcp_tools_result {
Ok(mcp_tools) => {
sess.late_mcp_tools.clear();
let auth = sess.services.auth_manager.auth().await;
connectors::refresh_accessible_connectors_cache_from_mcp_tools(
&turn_context.config,
codex-rs/core/src/plugins/injection.rs
+33 -17
View File
@@ -11,10 +11,17 @@ use crate::plugins::render_explicit_plugin_instructions;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_mcp::ToolInfo;
#[derive(Clone, Copy)]
pub(crate) enum PluginCapabilityResolution {
Runtime,
Declared,
}
pub(crate) fn build_plugin_injections(
mentioned_plugins: &[PluginCapabilitySummary],
mcp_tools: &[ToolInfo],
available_connectors: &[connectors::AppInfo],
resolution: PluginCapabilityResolution,
) -> Vec<ResponseItem> {
if mentioned_plugins.is_empty() {
return Vec::new();
@@ -25,27 +32,36 @@ pub(crate) fn build_plugin_injections(
mentioned_plugins
.iter()
.filter_map(|plugin| {
let available_mcp_servers = mcp_tools
.iter()
.filter(|tool| {
tool.server_name != CODEX_APPS_MCP_SERVER_NAME
&& tool
.plugin_display_names
.iter()
.any(|plugin_name| plugin_name == &plugin.display_name)
})
.map(|tool| tool.server_name.clone())
.collect::<BTreeSet<String>>()
.into_iter()
.collect::<Vec<_>>();
let available_mcp_servers = match resolution {
PluginCapabilityResolution::Runtime => mcp_tools
.iter()
.filter(|tool| {
tool.server_name != CODEX_APPS_MCP_SERVER_NAME
&& tool
.plugin_display_names
.iter()
.any(|plugin_name| plugin_name == &plugin.display_name)
})
.map(|tool| tool.server_name.clone())
.collect::<BTreeSet<String>>()
.into_iter()
.collect::<Vec<_>>(),
PluginCapabilityResolution::Declared => plugin.mcp_server_names.clone(),
};
let available_apps = available_connectors
.iter()
.filter(|connector| {
connector.is_enabled
&& connector
.plugin_display_names
.iter()
.any(|plugin_name| plugin_name == &plugin.display_name)
&& match resolution {
PluginCapabilityResolution::Runtime => connector
.plugin_display_names
.iter()
.any(|plugin_name| plugin_name == &plugin.display_name),
PluginCapabilityResolution::Declared => plugin
.app_connector_ids
.iter()
.any(|connector_id| connector_id.0 == connector.id),
}
})
.map(connector_display_label)
.collect::<BTreeSet<String>>()
codex-rs/core/src/plugins/mod.rs
+1
View File
@@ -8,6 +8,7 @@ pub(crate) mod test_support;
pub(crate) use codex_plugin::PluginCapabilitySummary;
pub(crate) use discoverable::list_tool_suggest_discoverable_plugins;
pub(crate) use injection::PluginCapabilityResolution;
pub(crate) use injection::build_plugin_injections;
pub(crate) use render::render_explicit_plugin_instructions;
codex-rs/core/src/session/mcp.rs
+1
View File
@@ -373,6 +373,7 @@ impl Session {
let mut manager = self.services.mcp_connection_manager.write().await;
std::mem::replace(&mut *manager, refreshed_manager)
};
self.late_mcp_tools.clear();
old_manager.shutdown().await;
}
codex-rs/core/src/session/mod.rs
+5 -1
View File
@@ -1522,6 +1522,7 @@ impl Session {
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
self.services.skills_manager.clear_cache();
self.services.plugins_manager.clear_cache();
self.late_mcp_tools.clear();
let hooks = build_hooks_for_config(
config.as_ref(),
self.services.plugins_manager.as_ref(),
@@ -2846,7 +2847,10 @@ impl Session {
.push(PersonalitySpecInstructions::new(personality_message).render());
}
}
if turn_context.config.include_apps_instructions && turn_context.apps_enabled() {
if turn_context.config.include_apps_instructions
&& turn_context.apps_enabled()
&& !crate::tools::lazy_mcp::enabled(turn_context)
{
let mcp_connection_manager = self.services.mcp_connection_manager.read().await;
let accessible_and_enabled_connectors =
connectors::list_accessible_and_enabled_connectors_from_manager(
codex-rs/core/src/session/session.rs
+3
View File
@@ -4,6 +4,7 @@ use crate::agents_md::LoadedAgentsMd;
use crate::config::ConstraintError;
use crate::skills::SkillError;
use crate::state::ActiveTurn;
use crate::tools::registry::LateToolRegistry;
use codex_protocol::SessionId;
use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE;
use codex_protocol::config_types::ServiceTier;
@@ -40,6 +41,7 @@ pub(crate) struct Session {
pub(super) features: ManagedFeatures,
pub(super) multi_agent_version: OnceLock<MultiAgentVersion>,
pub(super) pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
pub(crate) late_mcp_tools: LateToolRegistry,
pub(crate) conversation: Arc<RealtimeConversationManager>,
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
pub(crate) input_queue: InputQueue,
@@ -1068,6 +1070,7 @@ impl Session {
features: config.features.clone(),
multi_agent_version,
pending_mcp_server_refresh_config: Mutex::new(None),
late_mcp_tools: LateToolRegistry::default(),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: InputQueue::new(),
codex-rs/core/src/session/tests.rs
+6
View File
@@ -572,6 +572,7 @@ fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> T
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
@@ -1440,10 +1441,12 @@ async fn refresh_runtime_config_refreshes_hooks() -> anyhow::Result<()> {
};
assert!(session.hooks().preview_session_start(&request).is_empty());
let late_mcp_generation = session.late_mcp_tools.generation();
let next_config = load_latest_config_for_session(&session).await;
session.refresh_runtime_config(next_config).await;
assert_eq!(session.hooks().preview_session_start(&request).len(), 1);
assert_ne!(session.late_mcp_tools.generation(), late_mcp_generation);
Ok(())
}
@@ -4920,6 +4923,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
features: config.features.clone(),
multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
pending_mcp_server_refresh_config: Mutex::new(None),
late_mcp_tools: crate::tools::registry::LateToolRegistry::default(),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: super::input_queue::InputQueue::new(),
@@ -6986,6 +6990,7 @@ where
features: config.features.clone(),
multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
pending_mcp_server_refresh_config: Mutex::new(None),
late_mcp_tools: crate::tools::registry::LateToolRegistry::default(),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: super::input_queue::InputQueue::new(),
@@ -9429,6 +9434,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() {
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
);
let item = ResponseItem::CustomToolCall {
codex-rs/core/src/session/turn.rs
+55 -33
View File
@@ -34,6 +34,7 @@ use crate::mentions::build_skill_name_counts;
use crate::mentions::collect_explicit_app_ids;
use crate::mentions::collect_explicit_plugin_mentions;
use crate::mentions::collect_tool_mentions_from_messages;
use crate::plugins::PluginCapabilityResolution;
use crate::plugins::build_plugin_injections;
use crate::responses_retry::ResponsesStreamRequest;
use crate::responses_retry::handle_retryable_response_stream_error;
@@ -53,6 +54,7 @@ use crate::stream_events_utils::record_completed_response_item_with_finalized_fa
use crate::tasks::emit_compact_metric;
use crate::tools::ToolRouter;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::lazy_mcp;
use crate::tools::parallel::ToolCallRuntime;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::router::ToolRouterParams;
@@ -457,34 +459,41 @@ async fn build_skills_and_plugins(
// enabled plugins, then converted into turn-scoped guidance below.
let mentioned_plugins =
collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries());
let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() {
// Plugin mentions need raw MCP/app inventory even when app tools
// are normally hidden so we can describe the plugin's currently
// usable capabilities for this turn.
match sess
.services
.mcp_connection_manager
.read()
.await
.list_all_tools()
.or_cancel(cancellation_token)
.await
{
Ok(mcp_tools) => mcp_tools,
Err(_) if turn_context.apps_enabled() => return None,
Err(_) => Vec::new(),
}
} else {
Vec::new()
};
let lazy_mcp_enabled = lazy_mcp::enabled(turn_context);
let mcp_tools =
if !lazy_mcp_enabled && (turn_context.apps_enabled() || !mentioned_plugins.is_empty()) {
// Plugin mentions need raw MCP/app inventory even when app tools
// are normally hidden so we can describe the plugin's currently
// usable capabilities for this turn.
match sess
.services
.mcp_connection_manager
.read()
.await
.list_all_tools()
.or_cancel(cancellation_token)
.await
{
Ok(mcp_tools) => mcp_tools,
Err(_) if turn_context.apps_enabled() => return None,
Err(_) => Vec::new(),
}
} else {
Vec::new()
};
let available_connectors = if turn_context.apps_enabled() {
let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible(
loaded_plugins
.effective_apps()
.into_iter()
.map(|connector_id| connector_id.0),
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
);
let plugin_app_ids = loaded_plugins
.effective_apps()
.into_iter()
.map(|connector_id| connector_id.0);
let connectors = if lazy_mcp_enabled {
codex_connectors::merge::merge_plugin_connectors(Vec::new(), plugin_app_ids)
} else {
codex_connectors::merge::merge_plugin_connectors_with_accessible(
plugin_app_ids,
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
)
};
connectors::with_app_enabled_state(connectors, &turn_context.config)
} else {
Vec::new()
@@ -540,8 +549,16 @@ async fn build_skills_and_plugins(
&available_connectors,
&skill_name_counts_lower,
);
let plugin_items =
build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors);
let plugin_items = build_plugin_injections(
&mentioned_plugins,
&mcp_tools,
&available_connectors,
if lazy_mcp_enabled {
PluginCapabilityResolution::Declared
} else {
PluginCapabilityResolution::Runtime
},
);
let mut explicitly_enabled_connectors = collect_explicit_app_ids(&user_input);
explicitly_enabled_connectors.extend(skill_connector_ids);
let connector_names_by_id = available_connectors
@@ -1086,10 +1103,14 @@ pub(crate) async fn built_tools(
.instrument(trace_span!("read_mcp_connection_manager"))
.await;
let has_mcp_servers = mcp_connection_manager.has_servers();
let all_mcp_tools = mcp_connection_manager
.list_all_tools()
.or_cancel(cancellation_token)
.await?;
let all_mcp_tools = if lazy_mcp::enabled(turn_context) {
Vec::new()
} else {
mcp_connection_manager
.list_all_tools()
.or_cancel(cancellation_token)
.await?
};
drop(mcp_connection_manager);
let loaded_plugins = sess
.services
@@ -1171,6 +1192,7 @@ pub(crate) async fn built_tools(
discoverable_tools,
extension_tool_executors: extension_tool_executors(sess),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
late_mcp_tools: lazy_mcp::enabled(turn_context).then(|| sess.late_mcp_tools.clone()),
},
)))
}
codex-rs/core/src/stream_events_utils_tests.rs
+1
View File
@@ -273,6 +273,7 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() {
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
codex-rs/core/src/tools/handlers/request_plugin_install.rs
+1
View File
@@ -331,6 +331,7 @@ async fn refresh_missing_requested_connectors(
match manager.hard_refresh_codex_apps_tools_cache().await {
Ok(mcp_tools) => {
session.late_mcp_tools.clear();
let accessible_connectors = connectors::with_app_enabled_state(
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
&turn.config,
codex-rs/core/src/tools/handlers/tool_search.rs
+69 -11
View File
@@ -4,7 +4,9 @@ use crate::tools::context::ToolPayload;
use crate::tools::context::ToolSearchOutput;
use crate::tools::context::boxed_tool_output;
use crate::tools::handlers::tool_search_spec::create_tool_search_tool;
use crate::tools::lazy_mcp;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::LateToolRegistry;
use crate::tools::registry::ToolExecutor;
use bm25::Document;
use bm25::Language;
@@ -24,10 +26,19 @@ pub struct ToolSearchHandler {
entries: Vec<ToolSearchEntry>,
search_source_infos: Vec<ToolSearchSourceInfo>,
search_engine: SearchEngine<usize>,
late_mcp_tools: Option<LateToolRegistry>,
}
impl ToolSearchHandler {
#[cfg(test)]
pub(crate) fn new(search_infos: Vec<ToolSearchInfo>) -> Self {
Self::new_with_late_mcp_tools(search_infos, /*late_mcp_tools*/ None)
}
pub(crate) fn new_with_late_mcp_tools(
search_infos: Vec<ToolSearchInfo>,
late_mcp_tools: Option<LateToolRegistry>,
) -> Self {
let mut entries = Vec::with_capacity(search_infos.len());
let mut search_source_infos = Vec::new();
for search_info in search_infos {
@@ -36,19 +47,19 @@ impl ToolSearchHandler {
search_source_infos.push(source_info);
}
}
let documents: Vec<Document<usize>> = entries
.iter()
.map(|entry| entry.search_text.clone())
.enumerate()
.map(|(idx, search_text)| Document::new(idx, search_text))
.collect();
let search_engine =
SearchEngineBuilder::<usize>::with_documents(Language::English, documents).build();
if late_mcp_tools.is_some() {
search_source_infos.push(ToolSearchSourceInfo {
name: "MCP servers".to_string(),
description: Some("Tools provided by configured MCP servers.".to_string()),
});
}
let search_engine = Self::build_search_engine(&entries);
Self {
entries,
search_source_infos,
search_engine,
late_mcp_tools,
}
}
}
@@ -71,7 +82,12 @@ impl ToolExecutor<ToolInvocation> for ToolSearchHandler {
&self,
invocation: ToolInvocation,
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation { payload, .. } = invocation;
let ToolInvocation {
session,
turn,
payload,
..
} = invocation;
let args = match payload {
ToolPayload::ToolSearch { arguments } => arguments,
@@ -96,11 +112,27 @@ impl ToolExecutor<ToolInvocation> for ToolSearchHandler {
));
}
if self.entries.is_empty() {
let mut late_entries = Vec::new();
if let Some(late_mcp_tools) = &self.late_mcp_tools {
late_entries.extend(
lazy_mcp::register_eligible_tools(&session, &turn, late_mcp_tools)
.await
.into_iter()
.map(|search_info| search_info.entry),
);
}
if self.entries.is_empty() && late_entries.is_empty() {
return Ok(boxed_tool_output(ToolSearchOutput { tools: Vec::new() }));
}
let tools = self.search(query, limit)?;
let tools = if late_entries.is_empty() {
self.search(query, limit)?
} else {
let mut entries = self.entries.clone();
entries.extend(late_entries);
Self::search_entries(&entries, query, limit)?
};
Ok(boxed_tool_output(ToolSearchOutput { tools }))
}
@@ -109,6 +141,16 @@ impl ToolExecutor<ToolInvocation> for ToolSearchHandler {
impl CoreToolRuntime for ToolSearchHandler {}
impl ToolSearchHandler {
fn build_search_engine(entries: &[ToolSearchEntry]) -> SearchEngine<usize> {
let documents: Vec<Document<usize>> = entries
.iter()
.map(|entry| entry.search_text.clone())
.enumerate()
.map(|(idx, search_text)| Document::new(idx, search_text))
.collect();
SearchEngineBuilder::<usize>::with_documents(Language::English, documents).build()
}
fn search(
&self,
query: &str,
@@ -123,6 +165,22 @@ impl ToolSearchHandler {
self.search_output_tools(results)
}
fn search_entries(
entries: &[ToolSearchEntry],
query: &str,
limit: usize,
) -> Result<Vec<LoadableToolSpec>, FunctionCallError> {
let search_engine = Self::build_search_engine(entries);
let results = search_engine
.search(query, limit)
.into_iter()
.map(|result| result.document.id)
.filter_map(|id| entries.get(id));
Ok(coalesce_loadable_tool_specs(
results.map(|entry| entry.output.clone()),
))
}
fn search_output_tools<'a>(
&self,
results: impl IntoIterator<Item = &'a ToolSearchEntry>,
codex-rs/core/src/tools/lazy_mcp.rs
+86
View File
@@ -0,0 +1,86 @@
use crate::connectors;
use crate::mcp_tool_exposure::build_mcp_tool_exposure;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::tools::handlers::McpHandler;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::LateToolRegistry;
use crate::tools::registry::ToolExecutor;
use codex_features::Feature;
use codex_tools::ToolSearchInfo;
use std::sync::Arc;
use tracing::warn;
pub(crate) fn enabled(turn_context: &TurnContext) -> bool {
turn_context
.config
.features
.enabled(Feature::ToolSearchAlwaysDeferMcpTools)
&& turn_context.model_info.supports_search_tool
&& turn_context.provider.capabilities().namespace_tools
}
#[expect(
clippy::await_holding_invalid_type,
reason = "MCP tool listing borrows the read guard across the asynchronous inventory read"
)]
pub(crate) async fn register_eligible_tools(
session: &Arc<Session>,
turn_context: &TurnContext,
late_tools: &LateToolRegistry,
) -> Vec<ToolSearchInfo> {
let registry_generation = late_tools.generation();
let all_mcp_tools = session
.services
.mcp_connection_manager
.read()
.await
.list_all_tools()
.await;
let connectors = if turn_context.apps_enabled() {
let loaded_plugins = session
.services
.plugins_manager
.plugins_for_config(&turn_context.config.plugins_config_input())
.await;
let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible(
loaded_plugins
.effective_apps()
.into_iter()
.map(|connector_id| connector_id.0),
connectors::accessible_connectors_from_mcp_tools(&all_mcp_tools),
);
Some(connectors::with_app_enabled_state(
connectors,
&turn_context.config,
))
} else {
None
};
let exposure = build_mcp_tool_exposure(
&all_mcp_tools,
connectors.as_deref(),
&turn_context.config,
/*search_tool_enabled*/ true,
);
let mut search_infos = Vec::new();
let mut handlers: Vec<Arc<dyn CoreToolRuntime>> = Vec::new();
for tool_info in exposure.deferred_tools.unwrap_or(exposure.direct_tools) {
match McpHandler::new(tool_info) {
Ok(handler) => {
if let Some(search_info) = handler.search_info() {
search_infos.push(search_info);
}
handlers.push(Arc::new(handler));
}
Err(err) => warn!("Skipping deferred MCP tool: failed to build tool spec: {err}"),
}
}
if late_tools.replace_if_generation(registry_generation, handlers) {
search_infos
} else {
warn!("Skipping deferred MCP tools because the MCP inventory changed during tool_search");
Vec::new()
}
}
codex-rs/core/src/tools/mod.rs
+1
View File
@@ -4,6 +4,7 @@ pub(crate) mod events;
pub(crate) mod handlers;
pub(crate) mod hook_names;
pub(crate) mod hosted_spec;
pub(crate) mod lazy_mcp;
pub(crate) mod lifecycle;
pub(crate) mod network_approval;
pub(crate) mod orchestrator;
codex-rs/core/src/tools/registry.rs
+93 -3
View File
@@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::RwLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -324,11 +325,84 @@ impl CoreToolRuntime for ExposureOverride {
pub struct ToolRegistry {
tools: HashMap<ToolName, Arc<dyn CoreToolRuntime>>,
late_tools: Option<LateToolRegistry>,
}
#[derive(Clone, Default)]
pub(crate) struct LateToolRegistry {
state: Arc<RwLock<LateToolRegistryState>>,
}
#[derive(Default)]
struct LateToolRegistryState {
generation: u64,
tools: HashMap<ToolName, Arc<dyn CoreToolRuntime>>,
}
impl LateToolRegistry {
pub(crate) fn generation(&self) -> u64 {
self.state
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.generation
}
pub(crate) fn replace_if_generation(
&self,
generation: u64,
tools: impl IntoIterator<Item = Arc<dyn CoreToolRuntime>>,
) -> bool {
let tools = tools
.into_iter()
.map(|tool| (tool.tool_name(), tool))
.collect();
let mut state = self
.state
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.generation != generation {
return false;
}
state.tools = tools;
true
}
pub(crate) fn clear(&self) {
let mut state = self
.state
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.generation = state.generation.wrapping_add(1);
state.tools.clear();
}
fn tool(&self, name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
self.state
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.tools
.get(name)
.map(Arc::clone)
}
#[cfg(test)]
fn tool_names(&self) -> Vec<ToolName> {
self.state
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.tools
.keys()
.cloned()
.collect()
}
}
impl ToolRegistry {
fn new(tools: HashMap<ToolName, Arc<dyn CoreToolRuntime>>) -> Self {
Self { tools }
Self {
tools,
late_tools: None,
}
}
pub(crate) fn from_tools(tools: impl IntoIterator<Item = Arc<dyn CoreToolRuntime>>) -> Self {
@@ -344,6 +418,15 @@ impl ToolRegistry {
Self::new(tools_by_name)
}
pub(crate) fn from_tools_with_late_registry(
tools: impl IntoIterator<Item = Arc<dyn CoreToolRuntime>>,
late_tools: LateToolRegistry,
) -> Self {
let mut registry = Self::from_tools(tools);
registry.late_tools = Some(late_tools);
registry
}
#[cfg(test)]
pub(crate) fn empty_for_test() -> Self {
Self::new(HashMap::new())
@@ -359,19 +442,26 @@ impl ToolRegistry {
}
fn tool(&self, name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
self.tools.get(name).map(Arc::clone)
self.tools
.get(name)
.map(Arc::clone)
.or_else(|| self.late_tools.as_ref()?.tool(name))
}
#[cfg(test)]
pub(crate) fn tool_names_for_test(&self) -> Vec<ToolName> {
let mut names = self.tools.keys().cloned().collect::<Vec<_>>();
if let Some(late_tools) = &self.late_tools {
names.extend(late_tools.tool_names());
}
names.sort();
names.dedup();
names
}
#[cfg(test)]
pub(crate) fn tool_exposure(&self, name: &ToolName) -> Option<ToolExposure> {
self.tools.get(name).map(|tool| tool.exposure())
self.tool(name).map(|tool| tool.exposure())
}
pub(crate) fn create_diff_consumer(
codex-rs/core/src/tools/registry_tests.rs
+42
View File
@@ -27,6 +27,48 @@ impl ToolExecutor<ToolInvocation> for TestHandler {
impl CoreToolRuntime for TestHandler {}
#[test]
fn late_tool_registry_adds_tools_after_static_registry_creation() {
let static_name = codex_tools::ToolName::plain("static");
let late_name = codex_tools::ToolName::namespaced("mcp__late", "lookup");
let late_tools = LateToolRegistry::default();
let registry = ToolRegistry::from_tools_with_late_registry(
[Arc::new(TestHandler {
tool_name: static_name.clone(),
}) as Arc<dyn CoreToolRuntime>],
late_tools.clone(),
);
assert!(registry.tool(&static_name).is_some());
assert!(registry.tool(&late_name).is_none());
let generation = late_tools.generation();
assert!(late_tools.replace_if_generation(
generation,
[Arc::new(TestHandler {
tool_name: late_name.clone(),
}) as Arc<dyn CoreToolRuntime>],
));
assert!(registry.tool(&late_name).is_some());
assert_eq!(
registry.tool_names_for_test(),
vec![late_name.clone(), static_name.clone()]
);
late_tools.clear();
assert!(registry.tool(&late_name).is_none());
assert_eq!(registry.tool_names_for_test(), vec![static_name]);
assert!(!late_tools.replace_if_generation(
generation,
[Arc::new(TestHandler {
tool_name: late_name.clone(),
}) as Arc<dyn CoreToolRuntime>],
));
assert!(registry.tool(&late_name).is_none());
}
#[derive(Clone)]
enum LifecycleTestResult {
Ok { success: bool },
codex-rs/core/src/tools/router.rs
+2
View File
@@ -5,6 +5,7 @@ use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::registry::AnyToolResult;
use crate::tools::registry::LateToolRegistry;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::registry::ToolRegistry;
use crate::tools::spec_plan::build_tool_router;
@@ -42,6 +43,7 @@ pub(crate) struct ToolRouterParams<'a> {
pub(crate) discoverable_tools: Option<Vec<DiscoverableTool>>,
pub(crate) extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
pub(crate) dynamic_tools: &'a [DynamicToolSpec],
pub(crate) late_mcp_tools: Option<LateToolRegistry>,
}
impl ToolRouter {
codex-rs/core/src/tools/router_tests.rs
+5
View File
@@ -116,6 +116,7 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
);
@@ -195,6 +196,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
);
@@ -230,6 +232,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()>
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: turn.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
);
@@ -282,6 +285,7 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> {
discoverable_tools: None,
extension_tool_executors: Vec::new(),
dynamic_tools: &dynamic_tools,
late_mcp_tools: None,
},
);
@@ -343,6 +347,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow
discoverable_tools: None,
extension_tool_executors: extension_tool_executors(&session),
dynamic_tools: turn.dynamic_tools.as_slice(),
late_mcp_tools: None,
},
);
codex-rs/core/src/tools/spec_plan.rs
+17 -4
View File
@@ -48,6 +48,7 @@ use crate::tools::hosted_spec::WebSearchToolOptions;
use crate::tools::hosted_spec::create_image_generation_tool;
use crate::tools::hosted_spec::create_web_search_tool;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::LateToolRegistry;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::ToolRegistry;
use crate::tools::registry::override_tool_exposure;
@@ -146,6 +147,7 @@ struct CoreToolPlanContext<'a> {
dynamic_tools: &'a [DynamicToolSpec],
default_agent_type_description: &'a str,
wait_agent_timeouts: WaitAgentTimeoutOptions,
late_mcp_tools: Option<&'a LateToolRegistry>,
}
pub(crate) fn build_tool_router(
@@ -166,6 +168,7 @@ fn build_tool_specs_and_registry(
discoverable_tools,
extension_tool_executors,
dynamic_tools,
late_mcp_tools,
} = params;
let default_agent_type_description =
crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new());
@@ -178,17 +181,19 @@ fn build_tool_specs_and_registry(
dynamic_tools,
default_agent_type_description: &default_agent_type_description,
wait_agent_timeouts: wait_agent_timeout_options(turn_context),
late_mcp_tools: late_mcp_tools.as_ref(),
};
let mut planned_tools = PlannedTools::default();
add_tool_sources(&context, &mut planned_tools);
append_tool_search_executor(&context, &mut planned_tools);
prepend_code_mode_executors(&context, &mut planned_tools);
build_model_visible_specs_and_registry(turn_context, planned_tools)
build_model_visible_specs_and_registry(turn_context, planned_tools, late_mcp_tools)
}
fn build_model_visible_specs_and_registry(
turn_context: &TurnContext,
planned_tools: PlannedTools,
late_mcp_tools: Option<LateToolRegistry>,
) -> (Vec<ToolSpec>, ToolRegistry) {
let PlannedTools {
runtimes,
@@ -215,7 +220,12 @@ fn build_model_visible_specs_and_registry(
}
specs.extend(hosted_specs);
let registry = ToolRegistry::from_tools(runtimes);
let registry = match late_mcp_tools {
Some(late_mcp_tools) => {
ToolRegistry::from_tools_with_late_registry(runtimes, late_mcp_tools)
}
None => ToolRegistry::from_tools(runtimes),
};
let model_visible_specs = merge_into_namespaces(specs)
.into_iter()
.filter(|spec| {
@@ -852,11 +862,14 @@ fn append_tool_search_executor(
.filter(|executor| executor.exposure() == ToolExposure::Deferred)
.filter_map(|executor| executor.search_info())
.collect::<Vec<_>>();
if search_infos.is_empty() {
if search_infos.is_empty() && context.late_mcp_tools.is_none() {
return;
}
planned_tools.add(ToolSearchHandler::new(search_infos));
planned_tools.add(ToolSearchHandler::new_with_late_mcp_tools(
search_infos,
context.late_mcp_tools.cloned(),
));
}
fn prepend_code_mode_executors(
codex-rs/core/src/tools/spec_plan_tests.rs
+21
View File
@@ -33,6 +33,7 @@ use serde_json::json;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE;
use crate::tools::registry::LateToolRegistry;
use crate::tools::router::ToolRouter;
use crate::tools::router::ToolRouterParams;
@@ -43,6 +44,7 @@ struct ToolPlanInputs {
discoverable_tools: Option<Vec<DiscoverableTool>>,
extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
dynamic_tools: Vec<DynamicToolSpec>,
late_mcp_tools: Option<LateToolRegistry>,
}
struct ToolPlanProbe {
@@ -183,6 +185,7 @@ async fn probe_with(
discoverable_tools: inputs.discoverable_tools,
extension_tool_executors: inputs.extension_tool_executors,
dynamic_tools: inputs.dynamic_tools.as_slice(),
late_mcp_tools: inputs.late_mcp_tools,
},
);
ToolPlanProbe::from_router(router)
@@ -685,6 +688,24 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
"read_mcp_resource",
]);
let lazy_mcp = probe_with(
|turn| {
turn.model_info.supports_search_tool = true;
set_feature(
turn,
Feature::ToolSearchAlwaysDeferMcpTools,
/*enabled*/ true,
);
},
ToolPlanInputs {
late_mcp_tools: Some(LateToolRegistry::default()),
..ToolPlanInputs::default()
},
)
.await;
lazy_mcp.assert_visible_contains(&["tool_search"]);
lazy_mcp.assert_registered_contains(&["tool_search"]);
let bedrock_namespace_capability = probe_with(
|turn| {
turn.model_info.supports_search_tool = true;
codex-rs/core/tests/common/apps_test_server.rs
+139 -1
View File
@@ -7,6 +7,13 @@ use codex_login::CodexAuth;
use codex_models_manager::bundled_models_response;
use serde_json::Value;
use serde_json::json;
use std::sync::Arc;
use std::sync::Condvar;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::Notify;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::Request;
@@ -52,6 +59,69 @@ pub struct AppsTestServer {
pub chatgpt_base_url: String,
}
#[derive(Clone)]
pub struct ToolsListGate {
inner: Arc<ToolsListGateInner>,
}
struct ToolsListGateInner {
started: AtomicBool,
started_notify: Notify,
released: Mutex<bool>,
released_notify: Condvar,
}
impl ToolsListGate {
fn new() -> Self {
Self {
inner: Arc::new(ToolsListGateInner {
started: AtomicBool::new(false),
started_notify: Notify::new(),
released: Mutex::new(false),
released_notify: Condvar::new(),
}),
}
}
pub async fn wait_until_started(&self) {
while !self.inner.started.load(Ordering::Acquire) {
self.inner.started_notify.notified().await;
}
}
pub fn release(&self) {
*self
.inner
.released
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = true;
self.inner.released_notify.notify_all();
}
fn block(&self) {
self.inner.started.store(true, Ordering::Release);
self.inner.started_notify.notify_waiters();
let mut released = self
.inner
.released
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
while !*released {
released = self
.inner
.released_notify
.wait(released)
.unwrap_or_else(std::sync::PoisonError::into_inner);
}
}
}
impl Drop for ToolsListGate {
fn drop(&mut self) {
self.release();
}
}
#[derive(Clone, Copy)]
pub enum AppsTestToolLoading {
Direct,
@@ -63,7 +133,47 @@ impl AppsTestServer {
Self::mount_with_connector_name(server, CONNECTOR_NAME).await
}
pub async fn mount_with_tools_delay(
server: &MockServer,
tools_list_delay: Duration,
) -> Result<Self> {
mount_oauth_metadata(server).await;
mount_connectors_directory(server).await;
mount_streamable_http_json_rpc(
server,
CONNECTOR_NAME.to_string(),
CONNECTOR_DESCRIPTION.to_string(),
/*searchable*/ false,
/*include_app_only_tool*/ false,
tools_list_delay,
)
.await;
Ok(Self {
chatgpt_base_url: server.uri(),
})
}
pub async fn mount_with_tools_gate(server: &MockServer) -> Result<(Self, ToolsListGate)> {
let gate = ToolsListGate::new();
mount_oauth_metadata(server).await;
mount_connectors_directory(server).await;
mount_streamable_http_json_rpc_with_gate(server, gate.clone()).await;
Ok((
Self {
chatgpt_base_url: server.uri(),
},
gate,
))
}
pub async fn mount_searchable(server: &MockServer) -> Result<Self> {
Self::mount_searchable_with_tools_delay(server, Duration::ZERO).await
}
pub async fn mount_searchable_with_tools_delay(
server: &MockServer,
tools_list_delay: Duration,
) -> Result<Self> {
mount_oauth_metadata(server).await;
mount_connectors_directory(server).await;
mount_streamable_http_json_rpc(
@@ -72,6 +182,7 @@ impl AppsTestServer {
CONNECTOR_DESCRIPTION.to_string(),
/*searchable*/ true,
/*include_app_only_tool*/ false,
tools_list_delay,
)
.await;
Ok(Self {
@@ -91,6 +202,7 @@ impl AppsTestServer {
CONNECTOR_DESCRIPTION.to_string(),
/*searchable*/ false,
/*include_app_only_tool*/ false,
Duration::ZERO,
)
.await;
Ok(Self {
@@ -110,6 +222,7 @@ impl AppsTestServer {
CONNECTOR_DESCRIPTION.to_string(),
matches!(tool_loading, AppsTestToolLoading::Searchable),
/*include_app_only_tool*/ true,
Duration::ZERO,
)
.await;
Ok(Self {
@@ -264,6 +377,7 @@ async fn mount_streamable_http_json_rpc(
connector_description: String,
searchable: bool,
include_app_only_tool: bool,
tools_list_delay: Duration,
) {
Mock::given(method("POST"))
.and(path_regex("^/api/codex/apps/?$"))
@@ -272,6 +386,23 @@ async fn mount_streamable_http_json_rpc(
connector_description,
searchable,
include_app_only_tool,
tools_list_delay,
tools_list_gate: None,
})
.mount(server)
.await;
}
async fn mount_streamable_http_json_rpc_with_gate(server: &MockServer, gate: ToolsListGate) {
Mock::given(method("POST"))
.and(path_regex("^/api/codex/apps/?$"))
.respond_with(CodexAppsJsonRpcResponder {
connector_name: CONNECTOR_NAME.to_string(),
connector_description: CONNECTOR_DESCRIPTION.to_string(),
searchable: false,
include_app_only_tool: false,
tools_list_delay: Duration::ZERO,
tools_list_gate: Some(gate),
})
.mount(server)
.await;
@@ -282,6 +413,8 @@ struct CodexAppsJsonRpcResponder {
connector_description: String,
searchable: bool,
include_app_only_tool: bool,
tools_list_delay: Duration,
tools_list_gate: Option<ToolsListGate>,
}
impl Respond for CodexAppsJsonRpcResponder {
@@ -327,6 +460,9 @@ impl Respond for CodexAppsJsonRpcResponder {
}
"notifications/initialized" => ResponseTemplate::new(202),
"tools/list" => {
if let Some(gate) = &self.tools_list_gate {
gate.block();
}
let id = body.get("id").cloned().unwrap_or(Value::Null);
let mut response = json!({
"jsonrpc": "2.0",
@@ -475,7 +611,9 @@ impl Respond for CodexAppsJsonRpcResponder {
}
}));
}
ResponseTemplate::new(200).set_body_json(response)
ResponseTemplate::new(200)
.set_body_json(response)
.set_delay(self.tools_list_delay)
}
"tools/call" => {
let id = body.get("id").cloned().unwrap_or(Value::Null);
codex-rs/core/tests/suite/plugins.rs
+92
View File
@@ -11,6 +11,7 @@ use codex_login::CodexAuth;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use core_test_support::apps_test_server::AppsTestServer;
use core_test_support::apps_test_server::configure_search_capable_model;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::mount_sse_once;
@@ -23,6 +24,7 @@ use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use core_test_support::wait_for_mcp_server;
use tempfile::TempDir;
use tokio::time::timeout;
use wiremock::MockServer;
const SAMPLE_PLUGIN_CONFIG_NAME: &str = "sample@test";
@@ -326,6 +328,96 @@ async fn explicit_plugin_mentions_inject_plugin_guidance() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn lazy_explicit_plugin_mentions_use_declared_capabilities() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let apps_mock_server = MockServer::start().await;
let (apps_server, tools_list_gate) =
AppsTestServer::mount_with_tools_gate(&apps_mock_server).await?;
let mock = mount_sse_once(
&server,
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
)
.await;
let codex_home = Arc::new(TempDir::new()?);
let rmcp_test_server_bin = match stdio_server_bin() {
Ok(bin) => bin,
Err(err) => {
eprintln!("test_stdio_server binary not available, skipping test: {err}");
return Ok(());
}
};
write_plugin_skill_plugin(codex_home.as_ref());
write_plugin_mcp_plugin(codex_home.as_ref(), &rmcp_test_server_bin);
write_plugin_app_plugin(codex_home.as_ref());
let chatgpt_base_url = apps_server.chatgpt_base_url;
let mut builder = test_codex()
.with_home(codex_home)
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(move |config| {
config
.features
.enable(Feature::Apps)
.expect("test config should allow feature update");
config
.features
.enable(Feature::ToolSearchAlwaysDeferMcpTools)
.expect("test config should allow feature update");
configure_search_capable_model(config);
config.chatgpt_base_url = chatgpt_base_url;
});
let test = builder.build(&server).await?;
timeout(Duration::from_secs(/*secs*/ 10), async {
test.codex
.submit(Op::UserInput {
items: vec![codex_protocol::user_input::UserInput::Mention {
name: "sample".into(),
path: format!("plugin://{SAMPLE_PLUGIN_CONFIG_NAME}"),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
Result::<()>::Ok(())
})
.await
.expect("plugin mention should not wait for uncached MCP tools")?;
timeout(
Duration::from_secs(/*secs*/ 10),
tools_list_gate.wait_until_started(),
)
.await
.expect("background MCP tools/list should reach the deliberate gate");
tools_list_gate.release();
let request = mock.single_request();
let developer_messages = request.message_input_texts("developer");
assert!(
developer_messages
.iter()
.any(|text| text.contains("MCP servers from this plugin")),
"expected declared plugin MCP guidance: {developer_messages:?}"
);
assert!(
developer_messages
.iter()
.any(|text| text.contains("Apps from this plugin")),
"expected declared plugin app guidance: {developer_messages:?}"
);
let request_tools = tool_names(&request.body_json());
assert!(request_tools.iter().any(|name| name == "tool_search"));
assert!(!request_tools.iter().any(|name| name == "mcp__sample"));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn explicit_plugin_mentions_track_plugin_used_analytics() -> Result<()> {
skip_if_no_network!(Ok(()));
codex-rs/core/tests/suite/search_tool.rs
+21 -7
View File
@@ -53,6 +53,8 @@ use serde_json::Value;
use serde_json::json;
use std::collections::HashMap;
use std::time::Duration;
use tokio::time::timeout;
use wiremock::MockServer;
const SEARCH_TOOL_DESCRIPTION_SNIPPETS: [&str; 2] = [
"You have access to tools from the following sources",
@@ -180,7 +182,9 @@ async fn always_defer_feature_hides_small_app_tool_sets() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let apps_server = AppsTestServer::mount(&server).await?;
let apps_mock_server = MockServer::start().await;
let (apps_server, tools_list_gate) =
AppsTestServer::mount_with_tools_gate(&apps_mock_server).await?;
let mock = mount_sse_once(
&server,
sse(vec![
@@ -200,12 +204,23 @@ async fn always_defer_feature_hides_small_app_tool_sets() -> Result<()> {
});
let test = builder.build(&server).await?;
test.submit_turn_with_approval_and_permission_profile(
"list tools",
AskForApproval::Never,
PermissionProfile::Disabled,
timeout(
Duration::from_secs(/*secs*/ 10),
test.submit_turn_with_approval_and_permission_profile(
"list tools",
AskForApproval::Never,
PermissionProfile::Disabled,
),
)
.await?;
.await
.expect("first turn should not wait for uncached MCP tools")?;
timeout(
Duration::from_secs(/*secs*/ 10),
tools_list_gate.wait_until_started(),
)
.await
.expect("background MCP tools/list should reach the deliberate gate");
tools_list_gate.release();
let body = mock.single_request().body_json();
let tools = tool_names(&body);
@@ -1239,7 +1254,6 @@ async fn tool_search_surfaced_mcp_tool_errors_are_returned_to_model() -> Result<
.expect("test mcp servers should accept any configuration");
});
let test = builder.build(&server).await?;
wait_for_mcp_server(&test.codex, "rmcp").await?;
test.codex
.submit(Op::UserInput {