github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 08:01:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Keep effective tool exposure in the registry (#36367)

Browse Source 
## Why

`ToolExecutor::exposure()` describes a tool's preferred exposure, while the host
may apply step-specific policy when building a tool plan.

## What changed

- Store each runtime together with its effective exposure in `ToolRegistry`.
- Apply MCP, namespace, and collaboration-tool exposure overrides directly to
  registry entries instead of wrapping runtimes.
- Use the registry-owned exposure when building model-visible specs, deferred
  tool search, code-mode executors, and parallel-call support.

## Testing

- Cover that a hidden MCP tool remains routable but is not eligible for parallel
  tool calls.

GitOrigin-RevId: 1bc2760add8d993f477954269fa13681bfcc2444
This commit is contained in:
jif
2026-07-31 17:24:33 +00:00
committed by copyberry
parent 287e1020ae
commit 775fb21d2a
11 changed files with 235 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp_tool_exposure.rs
+10 -12
View File
@@ -13,16 +13,16 @@ use tracing::warn;
use crate::config::Config;
use crate::connectors;
use crate::tools::handlers::McpHandler;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::override_tool_exposure;
use crate::tools::registry::ToolRegistry;
#[instrument(level = "trace", skip_all)]
pub(crate) fn build_mcp_tool_runtimes<'a>(
all_mcp_tools: &'a [McpToolInfo],
connectors: Option<&'a [connectors::AppInfo]>,
config: &'a Config,
pub(crate) fn append_mcp_tools(
all_mcp_tools: &[McpToolInfo],
connectors: Option<&[connectors::AppInfo]>,
config: &Config,
search_tool_enabled: bool,
) -> impl Iterator<Item = Arc<dyn CoreToolRuntime>> + 'a {
registry: &mut ToolRegistry,
) {
// Keep regular MCP tools first; Apps tools also require connector and policy checks.
let non_app_tools = filter_non_codex_apps_mcp_tools_only(all_mcp_tools);
let app_tools = connectors
@@ -33,19 +33,17 @@ pub(crate) fn build_mcp_tool_runtimes<'a>(
} else {
ToolExposure::Direct
};
non_app_tools.chain(app_tools).filter_map(move |tool| {
for tool in non_app_tools.chain(app_tools) {
let tool_name = tool.canonical_tool_name();
match McpHandler::new(tool.clone()) {
Ok(handler) => {
let handler: Arc<dyn CoreToolRuntime> = Arc::new(handler);
Some(override_tool_exposure(handler, exposure))
registry.register_external_with_exposure(Arc::new(handler), exposure);
}
Err(err) => {
warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}");
None
}
}
})
}
}
fn filter_non_codex_apps_mcp_tools_only(
codex-rs/core/src/mcp_tool_exposure_test.rs
+26 -14
View File
@@ -12,6 +12,7 @@ use rmcp::model::Tool;
use super::*;
use crate::config::CONFIG_TOML_FILE;
use crate::config::Config;
use crate::config::ConfigBuilder;
use crate::config::test_config;
use crate::connectors::AppInfo;
@@ -91,11 +92,22 @@ fn expected_runtimes(
}
fn runtimes_by_name(
runtimes: impl IntoIterator<Item = Arc<dyn CoreToolRuntime>>,
tools: &[ToolInfo],
connectors: Option<&[AppInfo]>,
config: &Config,
search_tool_enabled: bool,
) -> HashMap<ToolName, ToolExposure> {
runtimes
.into_iter()
.map(|runtime| (runtime.tool_name(), runtime.exposure()))
let mut registry = ToolRegistry::default();
append_mcp_tools(
tools,
connectors,
config,
search_tool_enabled,
&mut registry,
);
registry
.entries()
.map(|tool| (tool.runtime.tool_name(), tool.exposure))
.collect()
}
@@ -114,12 +126,12 @@ async fn directly_exposes_effective_tool_sets_when_search_is_unavailable() {
let config = test_config().await;
let mcp_tools = numbered_mcp_tools(/*count*/ 2);
let runtimes = build_mcp_tool_runtimes(
let runtimes = runtimes_by_name(
&mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ false,
);
assert_eq!(
runtimes_by_name(runtimes),
runtimes,
expected_runtimes(&mcp_tools, ToolExposure::Direct)
);
}
@@ -188,7 +200,7 @@ async fn excludes_tools_hidden_from_model_exposure() {
];
let connectors = vec![make_connector("calendar", "Calendar")];
let runtimes = build_mcp_tool_runtimes(
let runtimes = runtimes_by_name(
&mcp_tools,
Some(connectors.as_slice()),
&config,
@@ -196,7 +208,7 @@ async fn excludes_tools_hidden_from_model_exposure() {
);
assert_eq!(
runtimes_by_name(runtimes),
runtimes,
expected_runtimes(&[visible_tool, visible_app_tool], ToolExposure::Direct)
);
}
@@ -239,7 +251,7 @@ enabled = true
let connectors = vec![make_connector("calendar", "Calendar")];
let mcp_tools = [enabled_tool.clone(), disabled_tool];
let runtimes = build_mcp_tool_runtimes(
let runtimes = runtimes_by_name(
&mcp_tools,
Some(connectors.as_slice()),
&config,
@@ -247,7 +259,7 @@ enabled = true
);
assert_eq!(
runtimes_by_name(runtimes),
runtimes,
expected_runtimes(&[enabled_tool], ToolExposure::Direct)
);
}
@@ -257,12 +269,12 @@ async fn defers_effective_tool_sets_when_search_is_available() {
let config = test_config().await;
let mcp_tools = numbered_mcp_tools(/*count*/ 2);
let runtimes = build_mcp_tool_runtimes(
let runtimes = runtimes_by_name(
&mcp_tools, /*connectors*/ None, &config, /*search_tool_enabled*/ true,
);
assert_eq!(
runtimes_by_name(runtimes),
runtimes,
expected_runtimes(&mcp_tools, ToolExposure::Deferred)
);
}
@@ -290,7 +302,7 @@ async fn defers_apps_and_non_app_mcp_tools() {
];
let connectors = vec![make_connector("calendar", "Calendar")];
let runtimes = build_mcp_tool_runtimes(
let runtimes = runtimes_by_name(
&mcp_tools,
Some(connectors.as_slice()),
&config,
@@ -298,7 +310,7 @@ async fn defers_apps_and_non_app_mcp_tools() {
);
assert_eq!(
runtimes_by_name(runtimes),
runtimes,
expected_runtimes(&mcp_tools, ToolExposure::Deferred)
);
}
codex-rs/core/src/session/tests.rs
+12 -13
View File
@@ -16,7 +16,6 @@ use crate::shell::default_user_shell;
use crate::shell_snapshot::ShellSnapshot;
use crate::test_support::models_manager_with_provider;
use crate::tools::format_exec_output_str;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::ToolRegistry;
use codex_config::ConfigLayerStack;
use codex_config::ConfigLayerStackOrdering;
@@ -780,32 +779,32 @@ async fn preview_session_start_hooks(
)
}
pub(crate) fn tool_runtimes_for_test_step(
pub(crate) fn tool_registry_for_test_step(
step_context: &StepContext,
) -> (Vec<Arc<dyn CoreToolRuntime>>, Vec<ToolSpec>) {
let mut tool_runtimes = crate::tools::spec_plan::build_core_tool_runtimes(
) -> (ToolRegistry, Vec<ToolSpec>) {
let mut registry = crate::tools::spec_plan::build_core_tool_registry(
step_context.turn.as_ref(),
&step_context.environments,
step_context.mcp.as_ref(),
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
);
let hosted_specs = crate::tools::spec_plan::append_source_tool_runtimes(
let hosted_specs = crate::tools::spec_plan::append_source_tools(
step_context.turn.as_ref(),
&mut tool_runtimes,
&mut registry,
Vec::new(),
Vec::new(),
&step_context.turn.dynamic_tools,
);
(tool_runtimes, hosted_specs)
(registry, hosted_specs)
}
fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> ToolCallRuntime {
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref());
let router = Arc::new(ToolRouter::from_tools(
let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref());
let router = Arc::new(ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
hosted_specs,
&Default::default(),
));
@@ -10881,10 +10880,10 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() {
async fn fatal_tool_error_stops_turn_and_reports_error() {
let (session, turn_context, _rx) = make_session_and_context_with_rx().await;
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref());
let router = ToolRouter::from_tools(
let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref());
let router = ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
hosted_specs,
&Default::default(),
);
codex-rs/core/src/stream_events_utils_tests.rs
+4 -4
View File
@@ -8,7 +8,7 @@ use super::last_assistant_message_from_item;
use super::response_item_may_include_external_context;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session::tests::tool_runtimes_for_test_step;
use crate::session::tests::tool_registry_for_test_step;
use crate::tools::ToolRouter;
use crate::tools::parallel::ToolCallRuntime;
use crate::turn_diff_tracker::TurnDiffTracker;
@@ -278,10 +278,10 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() {
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref());
let router = Arc::new(ToolRouter::from_tools(
let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref());
let router = Arc::new(ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
hosted_specs,
&Default::default(),
));
codex-rs/core/src/tools/registry.rs
+38 -98
View File
@@ -35,7 +35,6 @@ use codex_protocol::protocol::EventMsg;
use codex_rollout::state_db;
use codex_shell_command::parse_command::parse_shell_script;
use codex_tools::ToolName;
use codex_tools::ToolSearchInfo;
use codex_tools::ToolSpec;
use futures::future::BoxFuture;
use indexmap::IndexMap;
@@ -243,94 +242,15 @@ pub(crate) struct PostToolUsePayload {
pub(crate) tool_response: Value,
}
pub(crate) fn override_tool_exposure(
handler: Arc<dyn CoreToolRuntime>,
exposure: ToolExposure,
) -> Arc<dyn CoreToolRuntime> {
if handler.exposure() == exposure {
return handler;
}
Arc::new(ExposureOverride { handler, exposure })
}
struct ExposureOverride {
handler: Arc<dyn CoreToolRuntime>,
exposure: ToolExposure,
}
impl ToolExecutor<ToolInvocation> for ExposureOverride {
fn tool_name(&self) -> ToolName {
self.handler.tool_name()
}
fn spec(&self) -> ToolSpec {
self.handler.spec()
}
fn exposure(&self) -> ToolExposure {
self.exposure
}
fn supports_parallel_tool_calls(&self) -> bool {
self.exposure != ToolExposure::Hidden && self.handler.supports_parallel_tool_calls()
}
fn search_info(&self) -> Option<ToolSearchInfo> {
self.handler.search_info()
}
fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
self.handler.handle(invocation)
}
}
impl CoreToolRuntime for ExposureOverride {
fn wait_until_ready<'a>(&'a self, session: &'a Arc<Session>) -> Option<BoxFuture<'a, ()>> {
self.handler.wait_until_ready(session)
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
self.handler.matches_kind(payload)
}
fn waits_for_runtime_cancellation(&self) -> bool {
self.handler.waits_for_runtime_cancellation()
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
self.handler.pre_tool_use_payload(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
result: &dyn ToolOutput,
) -> Option<PostToolUsePayload> {
self.handler.post_tool_use_payload(invocation, result)
}
fn with_updated_hook_input(
&self,
invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
self.handler
.with_updated_hook_input(invocation, updated_input)
}
fn telemetry_tags(&self, invocation: &ToolInvocation) -> ToolTelemetryTags {
self.handler.telemetry_tags(invocation)
}
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
self.handler.create_diff_consumer()
}
/// A tool runtime together with its effective exposure for the current step.
pub(crate) struct RegisteredTool {
pub(crate) runtime: Arc<dyn CoreToolRuntime>,
pub(crate) exposure: ToolExposure,
}
#[derive(Default)]
pub struct ToolRegistry {
tools: IndexMap<ToolName, Arc<dyn CoreToolRuntime>>,
tools: IndexMap<ToolName, RegisteredTool>,
}
impl ToolRegistry {
@@ -358,13 +278,22 @@ impl ToolRegistry {
where
T: CoreToolRuntime + 'static,
{
self.register_trusted(override_tool_exposure(Arc::new(handler), exposure));
self.register_trusted_with_exposure(Arc::new(handler), exposure);
}
pub(crate) fn register_trusted(&mut self, runtime: Arc<dyn CoreToolRuntime>) {
let exposure = runtime.exposure();
self.register_trusted_with_exposure(runtime, exposure);
}
pub(crate) fn register_trusted_with_exposure(
&mut self,
runtime: Arc<dyn CoreToolRuntime>,
exposure: ToolExposure,
) {
match self.tools.entry(runtime.tool_name()) {
Entry::Vacant(entry) => {
entry.insert(runtime);
entry.insert(RegisteredTool { runtime, exposure });
}
Entry::Occupied(entry) => {
let tool_name = entry.key();
@@ -380,10 +309,21 @@ impl ToolRegistry {
return;
}
self.tools.shift_insert(0, tool_name, runtime);
let exposure = runtime.exposure();
self.tools
.shift_insert(0, tool_name, RegisteredTool { runtime, exposure });
}
pub(crate) fn register_external(&mut self, runtime: Arc<dyn CoreToolRuntime>) -> bool {
let exposure = runtime.exposure();
self.register_external_with_exposure(runtime, exposure)
}
pub(crate) fn register_external_with_exposure(
&mut self,
runtime: Arc<dyn CoreToolRuntime>,
exposure: ToolExposure,
) -> bool {
let tool_name = runtime.tool_name();
if tool_name.namespace.is_none() && tool_name.name == "shell_command" {
tracing::warn!(tool_name = %tool_name, "skipping external tool with reserved name");
@@ -392,7 +332,7 @@ impl ToolRegistry {
match self.tools.entry(tool_name) {
Entry::Vacant(entry) => {
entry.insert(runtime);
entry.insert(RegisteredTool { runtime, exposure });
true
}
Entry::Occupied(entry) => {
@@ -406,21 +346,21 @@ impl ToolRegistry {
}
pub(crate) fn remove(&mut self, tool_name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
self.tools.shift_remove(tool_name)
self.tools.shift_remove(tool_name).map(|tool| tool.runtime)
}
pub(crate) fn runtimes(&self) -> impl Iterator<Item = &Arc<dyn CoreToolRuntime>> {
pub(crate) fn entries(&self) -> impl Iterator<Item = &RegisteredTool> {
self.tools.values()
}
pub(crate) fn runtimes_mut(&mut self) -> impl Iterator<Item = &mut Arc<dyn CoreToolRuntime>> {
pub(crate) fn entries_mut(&mut self) -> impl Iterator<Item = &mut RegisteredTool> {
self.tools.values_mut()
}
pub(crate) fn deferred_tool_namespaces(&self) -> BTreeMap<String, String> {
let mut namespaces = BTreeMap::<String, String>::new();
for (name, tool) in &self.tools {
if tool.exposure() != ToolExposure::Deferred {
if tool.exposure != ToolExposure::Deferred {
continue;
}
let Some(namespace) = &name.namespace else {
@@ -430,7 +370,7 @@ impl ToolRegistry {
if !existing_description.trim().is_empty() {
continue;
}
let description = match tool.spec() {
let description = match tool.runtime.spec() {
ToolSpec::Namespace(namespace) => namespace.description,
ToolSpec::Function(_)
| ToolSpec::Freeform(_)
@@ -458,7 +398,7 @@ impl ToolRegistry {
}
pub(crate) fn tool(&self, name: &ToolName) -> Option<Arc<dyn CoreToolRuntime>> {
self.tools.get(name).map(Arc::clone)
self.tools.get(name).map(|tool| Arc::clone(&tool.runtime))
}
#[cfg(test)]
@@ -470,7 +410,7 @@ impl ToolRegistry {
#[cfg(test)]
pub(crate) fn tool_exposure(&self, name: &ToolName) -> Option<ToolExposure> {
self.tools.get(name).map(|tool| tool.exposure())
self.tools.get(name).map(|tool| tool.exposure)
}
pub(crate) fn create_diff_consumer(
@@ -481,8 +421,8 @@ impl ToolRegistry {
}
pub(crate) fn supports_parallel_tool_calls(&self, name: &ToolName) -> Option<bool> {
let tool = self.tool(name)?;
Some(tool.supports_parallel_tool_calls())
let tool = self.tools.get(name)?;
Some(tool.exposure != ToolExposure::Hidden && tool.runtime.supports_parallel_tool_calls())
}
pub(crate) fn waits_for_runtime_cancellation(&self, name: &ToolName) -> Option<bool> {
codex-rs/core/src/tools/registry_tests.rs
+15 -13
View File
@@ -222,8 +222,8 @@ fn registry_preserves_external_winners_and_trusted_synthetic_order() {
assert_eq!(
registry
.runtimes()
.map(|runtime| runtime.tool_name())
.entries()
.map(|tool| tool.runtime.tool_name())
.collect::<Vec<_>>(),
vec![synthetic_name, first_name.clone(), second_name],
);
@@ -243,6 +243,10 @@ fn reserved_shell_command_rejects_external_runtimes_without_a_builtin() {
let mut registry = ToolRegistry::default();
assert!(!registry.register_external(handler(shell_command_name.clone())));
assert!(!registry.register_external_with_exposure(
handler(shell_command_name.clone()),
ToolExposure::Direct,
));
assert!(registry.tool(&shell_command_name).is_none());
let namespaced_handler = handler(namespaced_shell_command_name.clone());
@@ -255,7 +259,7 @@ fn reserved_shell_command_rejects_external_runtimes_without_a_builtin() {
}
#[tokio::test]
async fn readiness_selects_exact_tool_and_forwards_exposure_overrides() {
async fn readiness_selects_exact_tool_with_registry_owned_exposure() {
let (session, _turn) = crate::session::tests::make_session_and_context().await;
let session = Arc::new(session);
let plain_name = codex_tools::ToolName::plain("echo");
@@ -275,16 +279,14 @@ async fn readiness_selects_exact_tool_and_forwards_exposure_overrides() {
},
readiness_waits: Arc::clone(&plain_readiness_waits),
}) as Arc<dyn CoreToolRuntime>;
let namespaced_handler = override_tool_exposure(
Arc::new(ReadinessTestHandler {
handler: TestHandler {
tool_name: namespaced_name.clone(),
},
readiness_waits: Arc::clone(&namespaced_readiness_waits),
}),
ToolExposure::DirectModelOnly,
);
let registry = ToolRegistry::from_tools([plain_handler, namespaced_handler]);
let namespaced_handler = Arc::new(ReadinessTestHandler {
handler: TestHandler {
tool_name: namespaced_name.clone(),
},
readiness_waits: Arc::clone(&namespaced_readiness_waits),
});
let mut registry = ToolRegistry::from_tools([plain_handler]);
registry.register_trusted_with_exposure(namespaced_handler, ToolExposure::DirectModelOnly);
registry
.tool(&plain_name)
codex-rs/core/src/tools/router.rs
+3 -3
View File
@@ -84,15 +84,15 @@ pub(crate) struct ToolSuggestCandidates {
impl ToolRouter {
#[cfg(test)]
pub(crate) fn from_tools(
pub(crate) fn from_registry(
turn_context: &TurnContext,
tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
registry: ToolRegistry,
hosted_specs: Vec<ToolSpec>,
tool_search_handler_cache: &ToolSearchHandlerCache,
) -> Self {
finalize_tool_router(
turn_context,
ToolRegistry::from_tools(tool_runtimes),
registry,
hosted_specs,
tool_search_handler_cache,
)
codex-rs/core/src/tools/router_tests.rs
+42 -17
View File
@@ -7,10 +7,10 @@ use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::McpHandler;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::RegisteredTool;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::override_tool_exposure;
use crate::tools::spec_plan::append_source_tool_runtimes;
use crate::tools::spec_plan::build_core_tool_runtimes;
use crate::tools::spec_plan::append_source_tools;
use crate::tools::spec_plan::build_core_tool_registry;
use crate::tools::spec_plan::extension_tool_executors;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_extension_api::ExtensionData;
@@ -128,27 +128,27 @@ fn extension_tool_test_registry() -> Arc<ExtensionRegistry<Config>> {
fn test_tool_router(
step_context: &StepContext,
mcp_tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
mcp_tools: Vec<RegisteredTool>,
extension_tool_executors: impl IntoIterator<Item = Arc<dyn ToolExecutor<ExtensionToolCall>>>,
dynamic_tools: &[DynamicToolSpec],
) -> ToolRouter {
let mut tool_runtimes = build_core_tool_runtimes(
let mut registry = build_core_tool_registry(
step_context.turn.as_ref(),
&step_context.environments,
step_context.mcp.as_ref(),
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
);
let hosted_specs = append_source_tool_runtimes(
let hosted_specs = append_source_tools(
step_context.turn.as_ref(),
&mut tool_runtimes,
mcp_tool_runtimes,
&mut registry,
mcp_tools,
extension_tool_executors,
dynamic_tools,
);
ToolRouter::from_tools(
ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
hosted_specs,
&Default::default(),
)
@@ -282,15 +282,24 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
"mcp__echo__",
"query_with_delay",
)),
override_tool_exposure(
mcp_runtime(mcp_tool_info(
RegisteredTool {
exposure: ToolExposure::DirectModelOnly,
..mcp_runtime(mcp_tool_info(
"hello_echo",
/*supports_parallel_tool_calls*/ false,
"mcp__hello_echo__",
"query_with_delay",
)),
ToolExposure::DirectModelOnly,
),
))
},
RegisteredTool {
exposure: ToolExposure::Hidden,
..mcp_runtime(mcp_tool_info(
"hidden_echo",
/*supports_parallel_tool_calls*/ true,
"mcp__hidden_echo__",
"query_with_delay",
))
},
],
Vec::new(),
&turn.dynamic_tools,
@@ -328,6 +337,17 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> {
Some(different_server_call.tool_name.clone())
);
let hidden_call = ToolCall {
tool_name: ToolName::namespaced("mcp__hidden_echo__", "query_with_delay"),
call_id: "call-hidden".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
encrypted_function_args: None,
};
assert!(!router.tool_supports_parallel(&hidden_call));
assert!(router.tool_runtime(&hidden_call).is_some());
Ok(())
}
@@ -435,8 +455,13 @@ fn mcp_tool_info(
}
}
fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> Arc<dyn CoreToolRuntime> {
Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build"))
fn mcp_runtime(tool_info: codex_mcp::ToolInfo) -> RegisteredTool {
let runtime = Arc::new(McpHandler::new(tool_info).expect("MCP tool spec should build"))
as Arc<dyn CoreToolRuntime>;
RegisteredTool {
exposure: runtime.exposure(),
runtime,
}
}
#[tokio::test]
codex-rs/core/src/tools/spec_plan.rs
+48 -49
View File
@@ -2,7 +2,7 @@ use crate::agent::exceeds_thread_spawn_depth_limit;
use crate::agent::next_thread_spawn_depth;
use crate::connectors::AppInfo;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::mcp_tool_exposure::build_mcp_tool_runtimes;
use crate::mcp_tool_exposure::append_mcp_tools;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::tools::code_mode::default_exec_yield_time_override_ms;
@@ -56,9 +56,10 @@ 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::registry::CoreToolRuntime;
#[cfg(test)]
use crate::tools::registry::RegisteredTool;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::ToolRegistry;
use crate::tools::registry::override_tool_exposure;
use crate::tools::router::ToolRouter;
use codex_extension_api::ExtensionData;
use codex_features::Feature;
@@ -141,14 +142,13 @@ pub(crate) fn build_tool_router(
{
Vec::new()
} else {
for runtime in build_mcp_tool_runtimes(
append_mcp_tools(
mcp.tools(),
connectors,
&turn_context.config,
search_tool_enabled(turn_context),
) {
registry.register_external(runtime);
}
&mut registry,
);
let standalone_web_search_tool = append_extension_tool_executors(
turn_context,
extension_tool_executors(session, step_store),
@@ -168,13 +168,13 @@ pub(crate) fn build_tool_router(
#[cfg(test)]
#[instrument(level = "trace", skip_all)]
pub(crate) fn build_core_tool_runtimes(
pub(crate) fn build_core_tool_registry(
turn_context: &TurnContext,
environments: &TurnEnvironmentSnapshot,
mcp: &codex_mcp::McpBinding,
tool_suggest_candidates: Option<&crate::tools::router::ToolSuggestCandidates>,
wait_for_environment_tool_config: Option<&Arc<crate::WaitForEnvironmentToolConfig>>,
) -> Vec<Arc<dyn CoreToolRuntime>> {
) -> ToolRegistry {
let default_agent_type_description =
crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new());
let context = CoreToolPlanContext {
@@ -188,15 +188,15 @@ pub(crate) fn build_core_tool_runtimes(
};
let mut registry = ToolRegistry::default();
add_core_tool_sources(&context, &mut registry);
registry.runtimes().cloned().collect()
registry
}
#[cfg(test)]
#[instrument(level = "trace", skip_all)]
pub(crate) fn append_source_tool_runtimes(
pub(crate) fn append_source_tools(
turn_context: &TurnContext,
tool_runtimes: &mut Vec<Arc<dyn CoreToolRuntime>>,
mcp_tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
registry: &mut ToolRegistry,
mcp_tools: Vec<RegisteredTool>,
extension_tool_executors: impl IntoIterator<Item = Arc<dyn ToolExecutor<ExtensionToolCall>>>,
dynamic_tools: &[DynamicToolSpec],
) -> Vec<ToolSpec> {
@@ -204,14 +204,12 @@ pub(crate) fn append_source_tool_runtimes(
return Vec::new();
}
let mut registry = ToolRegistry::from_tools(std::mem::take(tool_runtimes));
for runtime in mcp_tool_runtimes {
registry.register_external(runtime);
for tool in mcp_tools {
registry.register_external_with_exposure(tool.runtime, tool.exposure);
}
let standalone_web_search_tool =
append_extension_tool_executors(turn_context, extension_tool_executors, &mut registry);
append_dynamic_tool_runtimes(dynamic_tools, &mut registry);
tool_runtimes.extend(registry.runtimes().cloned());
append_extension_tool_executors(turn_context, extension_tool_executors, registry);
append_dynamic_tool_runtimes(dynamic_tools, registry);
hosted_model_tool_specs(turn_context, standalone_web_search_tool.as_slice())
}
@@ -252,10 +250,10 @@ pub(crate) fn finalize_tool_router(
}
let tool_search_name = ToolName::plain(TOOL_SEARCH_TOOL_NAME);
if search_tool_enabled(turn_context)
&& registry.runtimes().any(|runtime| {
runtime.tool_name() != tool_search_name
&& runtime.exposure() == ToolExposure::Deferred
&& runtime.search_info().is_some()
&& registry.entries().any(|tool| {
tool.runtime.tool_name() != tool_search_name
&& tool.exposure == ToolExposure::Deferred
&& tool.runtime.search_info().is_some()
})
{
registry.remove(&tool_search_name);
@@ -273,8 +271,9 @@ fn apply_direct_model_only_namespace_overrides(
turn_context: &TurnContext,
registry: &mut ToolRegistry,
) {
for runtime in registry.runtimes_mut() {
let configured = runtime
for tool in registry.entries_mut() {
let configured = tool
.runtime
.tool_name()
.namespace
.as_ref()
@@ -285,8 +284,8 @@ fn apply_direct_model_only_namespace_overrides(
.direct_only_tool_namespaces
.contains(namespace)
});
if configured && runtime.exposure().is_available_in_code_mode() {
*runtime = override_tool_exposure(Arc::clone(runtime), ToolExposure::DirectModelOnly);
if configured && tool.exposure.is_available_in_code_mode() {
tool.exposure = ToolExposure::DirectModelOnly;
}
}
}
@@ -299,12 +298,12 @@ fn build_model_visible_specs(
hosted_specs: Vec<ToolSpec>,
) -> Vec<ToolSpec> {
let mut specs = Vec::new();
for runtime in registry.runtimes() {
let tool_name = runtime.tool_name();
let exposure = runtime.exposure();
for tool in registry.entries() {
let tool_name = tool.runtime.tool_name();
let exposure = tool.exposure;
if exposure.is_direct() && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure)
{
let spec = runtime.spec();
let spec = tool.runtime.spec();
specs.push(spec_for_model_request(
turn_context,
exposure,
@@ -519,18 +518,18 @@ fn register_code_mode_executors(
let mut exec_prompt_tool_specs = Vec::new();
let mut deferred_exec_prompt_tool_specs = Vec::new();
let deferred_tools_guidance_enabled = search_tool_enabled(turn_context);
for executor in registry.runtimes() {
let exposure = executor.exposure();
for tool in registry.entries() {
let exposure = tool.exposure;
if !exposure.is_available_in_code_mode() {
continue;
}
let tool_name = executor.tool_name();
let tool_name = tool.runtime.tool_name();
if is_excluded_from_code_mode(turn_context, &tool_name) {
continue;
}
let spec = executor.spec();
let spec = tool.runtime.spec();
let Some(code_mode_tool) = codex_tools::tool_spec_to_code_mode_tool_definition(&spec)
else {
continue;
@@ -900,7 +899,7 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut Too
agent_type_description(turn_context, context.default_agent_type_description);
let hide_spawn_agent_metadata =
turn_context.config.multi_agent_v2.hide_spawn_agent_metadata;
registry.register_trusted(override_tool_exposure(
registry.register_trusted_with_exposure(
multi_agent_v2_handler(
SpawnAgentHandlerV2::new(SpawnAgentToolOptions {
available_models: turn_context.available_models.clone(),
@@ -917,32 +916,32 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut Too
tool_namespace,
),
exposure,
));
registry.register_trusted(override_tool_exposure(
);
registry.register_trusted_with_exposure(
multi_agent_v2_handler(SendMessageHandlerV2, tool_namespace),
exposure,
));
registry.register_trusted(override_tool_exposure(
);
registry.register_trusted_with_exposure(
multi_agent_v2_handler(FollowupTaskHandlerV2, tool_namespace),
exposure,
));
);
if turn_context.config.multi_agent_v2.wait_agent_enabled {
registry.register_trusted(override_tool_exposure(
registry.register_trusted_with_exposure(
multi_agent_v2_handler(
WaitAgentHandlerV2::new(context.wait_agent_timeouts),
tool_namespace,
),
exposure,
));
);
}
registry.register_trusted(override_tool_exposure(
registry.register_trusted_with_exposure(
multi_agent_v2_handler(InterruptAgentHandler, tool_namespace),
exposure,
));
registry.register_trusted(override_tool_exposure(
);
registry.register_trusted_with_exposure(
multi_agent_v2_handler(ListAgentsHandlerV2, tool_namespace),
exposure,
));
);
} else {
let agent_type_description =
agent_type_description(turn_context, context.default_agent_type_description);
@@ -1012,9 +1011,9 @@ fn append_tool_search_executor(
tool_search_handler_cache: &ToolSearchHandlerCache,
) {
let search_infos = registry
.runtimes()
.filter(|executor| executor.exposure() == ToolExposure::Deferred)
.filter_map(|executor| executor.search_info())
.entries()
.filter(|tool| tool.exposure == ToolExposure::Deferred)
.filter_map(|tool| tool.runtime.search_info())
.collect::<Vec<_>>();
let source_listing = if turn_context
.config
codex-rs/core/src/tools/spec_plan_tests.rs
+36 -46
View File
@@ -41,18 +41,18 @@ use crate::tools::handlers::ToolSearchHandlerCache;
use crate::tools::handlers::WaitForEnvironmentHandler;
use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::override_tool_exposure;
use crate::tools::registry::RegisteredTool;
use crate::tools::router::ToolRouter;
use crate::tools::router::ToolSuggestCandidates;
use crate::tools::router::ToolSuggestPresentation;
use crate::tools::spec_plan::append_source_tool_runtimes;
use crate::tools::spec_plan::build_core_tool_runtimes;
use crate::tools::spec_plan::append_source_tools;
use crate::tools::spec_plan::build_core_tool_registry;
const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration";
#[derive(Default)]
struct ToolPlanInputs {
tool_runtimes: Vec<Arc<dyn CoreToolRuntime>>,
tool_runtimes: Vec<RegisteredTool>,
tool_suggest_candidates: Option<ToolSuggestCandidates>,
extension_tool_executors: Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>>,
wait_for_environment_tool_config: Option<Arc<WaitForEnvironmentToolConfig>>,
@@ -190,23 +190,23 @@ async fn probe_with(
configure_turn(&mut turn);
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let mut tool_runtimes = build_core_tool_runtimes(
let mut registry = build_core_tool_registry(
step_context.turn.as_ref(),
&step_context.environments,
step_context.mcp.as_ref(),
inputs.tool_suggest_candidates.as_ref(),
inputs.wait_for_environment_tool_config.as_ref(),
);
let hosted_specs = append_source_tool_runtimes(
let hosted_specs = append_source_tools(
step_context.turn.as_ref(),
&mut tool_runtimes,
&mut registry,
inputs.tool_runtimes,
inputs.extension_tool_executors,
&inputs.dynamic_tools,
);
let router = ToolRouter::from_tools(
let router = ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
hosted_specs,
&Default::default(),
);
@@ -396,11 +396,14 @@ fn mcp_runtime(
namespace: &str,
name: &str,
exposure: ToolExposure,
) -> Arc<dyn CoreToolRuntime> {
) -> RegisteredTool {
let handler: Arc<dyn CoreToolRuntime> = Arc::new(
McpHandler::new(mcp_tool(server, namespace, name)).expect("MCP tool spec should build"),
);
override_tool_exposure(handler, exposure)
RegisteredTool {
runtime: handler,
exposure,
}
}
fn dynamic_tool(namespace: Option<&str>, name: &str, defer_loading: bool) -> DynamicToolSpec {
@@ -914,9 +917,9 @@ async fn environment_tools_follow_the_step_context() {
&turn.config,
)));
let plan = ToolPlanProbe::from_router(ToolRouter::from_tools(
let plan = ToolPlanProbe::from_router(ToolRouter::from_registry(
turn.as_ref(),
build_core_tool_runtimes(
build_core_tool_registry(
turn.as_ref(),
&environments,
mcp.as_ref(),
@@ -1023,7 +1026,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
&["lookup".to_string()]
);
let searchable_mcp = ToolPlanInputs {
let searchable_mcp = || ToolPlanInputs {
tool_runtimes: vec![mcp_runtime(
"searchable",
"mcp__searchable",
@@ -1037,10 +1040,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
|turn| {
turn.model_info.supports_search_tool = false;
},
ToolPlanInputs {
tool_runtimes: searchable_mcp.tool_runtimes.clone(),
..ToolPlanInputs::default()
},
searchable_mcp(),
)
.await;
missing_model_capability.assert_visible_lacks(&["tool_search"]);
@@ -1062,10 +1062,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
turn.model_info.supports_search_tool = true;
use_bedrock_provider(turn);
},
ToolPlanInputs {
tool_runtimes: searchable_mcp.tool_runtimes.clone(),
..ToolPlanInputs::default()
},
searchable_mcp(),
)
.await;
bedrock_namespace_capability.assert_visible_contains(&["tool_search"]);
@@ -1074,7 +1071,7 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
|turn| {
turn.model_info.supports_search_tool = true;
},
searchable_mcp,
searchable_mcp(),
)
.await;
enabled.assert_visible_contains(&["tool_search"]);
@@ -1270,22 +1267,18 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
first_turn.model_info.supports_search_tool = true;
let first_turn = Arc::new(first_turn);
let first_step_context = StepContext::for_test(Arc::clone(&first_turn));
let mut first_tool_runtimes = build_core_tool_runtimes(
let mut first_registry = build_core_tool_registry(
first_step_context.turn.as_ref(),
&first_step_context.environments,
first_step_context.mcp.as_ref(),
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
);
first_tool_runtimes.push(mcp_runtime(
"first",
"mcp__first",
"lookup",
ToolExposure::Deferred,
));
let first_router = ToolRouter::from_tools(
let first_tool = mcp_runtime("first", "mcp__first", "lookup", ToolExposure::Deferred);
first_registry.register_external_with_exposure(first_tool.runtime, first_tool.exposure);
let first_router = ToolRouter::from_registry(
first_step_context.turn.as_ref(),
first_tool_runtimes,
first_registry,
super::hosted_model_tool_specs(first_step_context.turn.as_ref(), &[]),
&cache,
);
@@ -1295,22 +1288,18 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
second_turn.model_info.supports_search_tool = true;
let second_turn = Arc::new(second_turn);
let second_step_context = StepContext::for_test(Arc::clone(&second_turn));
let mut second_tool_runtimes = build_core_tool_runtimes(
let mut second_registry = build_core_tool_registry(
second_step_context.turn.as_ref(),
&second_step_context.environments,
second_step_context.mcp.as_ref(),
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
);
second_tool_runtimes.push(mcp_runtime(
"second",
"mcp__second",
"lookup",
ToolExposure::Deferred,
));
let second_router = ToolRouter::from_tools(
let second_tool = mcp_runtime("second", "mcp__second", "lookup", ToolExposure::Deferred);
second_registry.register_external_with_exposure(second_tool.runtime, second_tool.exposure);
let second_router = ToolRouter::from_registry(
second_step_context.turn.as_ref(),
second_tool_runtimes,
second_registry,
super::hosted_model_tool_specs(second_step_context.turn.as_ref(), &[]),
&cache,
);
@@ -1351,22 +1340,23 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() {
);
let turn = Arc::new(turn);
let step_context = StepContext::for_test(Arc::clone(&turn));
let mut tool_runtimes = build_core_tool_runtimes(
let mut registry = build_core_tool_registry(
step_context.turn.as_ref(),
&step_context.environments,
step_context.mcp.as_ref(),
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
);
tool_runtimes.push(mcp_runtime(
let tool = mcp_runtime(
"calendar",
"mcp__calendar",
"lookup",
ToolExposure::Deferred,
));
let router = ToolRouter::from_tools(
);
registry.register_external_with_exposure(tool.runtime, tool.exposure);
let router = ToolRouter::from_registry(
step_context.turn.as_ref(),
tool_runtimes,
registry,
super::hosted_model_tool_specs(step_context.turn.as_ref(), &[]),
&cache,
);
codex-rs/tools/src/tool_executor.rs
+1
View File
@@ -60,6 +60,7 @@ pub trait ToolExecutor<Invocation>: Send + Sync {
fn spec(&self) -> ToolSpec;
/// The preferred exposure before the host applies step-specific policy.
fn exposure(&self) -> ToolExposure {
ToolExposure::Direct
}