Introduce tool exposure for deferred registration (#22489)

## Why

Deferred tools were tracked with separate side-channel filtering after
tool specs had already been assembled. That made the registry
responsible for executing tools while the router/spec planner separately
decided whether those same tools should be exposed to the model up
front.

This PR makes exposure part of the tool handler contract so direct
versus deferred availability travels with the executable tool
registration.

Next step will be to simplify registration

## What Changed

- Adds `ToolExposure` to `codex-tools` and exposes it through
`ToolExecutor`, defaulting tools to `Direct`.
- Teaches dynamic tools and MCP handlers to mark deferred tools as
`Deferred` at construction time.
- Renames the registry object-safe wrapper from `AnyToolHandler` to
`RegisteredTool` and uses `ToolExposure` when deciding whether to
include a handler's spec in the initial model-visible tool list.
- Refactors tool spec planning to derive direct specs and deferred
search entries from registered handlers, removing the router's
special-case deferred dynamic tool filtering.

## Verification

- Not run.
This commit is contained in:
jif-oai
2026-05-13 18:16:51 +02:00
committed by GitHub
parent 889ee018e7
commit fdda59c00b
8 changed files with 100 additions and 105 deletions

View File

@@ -6,6 +6,7 @@ use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::parse_arguments;
use crate::tools::registry::ToolExecutor;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::ToolHandler;
use crate::tools::tool_search_entry::ToolSearchInfo;
use crate::turn_timing::now_unix_timestamp_ms;
@@ -30,6 +31,7 @@ use tracing::warn;
pub struct DynamicToolHandler {
tool_name: ToolName,
spec: Option<ToolSpec>,
exposure: ToolExposure,
search_text: String,
}
@@ -48,6 +50,11 @@ impl DynamicToolHandler {
Some(Self {
tool_name,
spec: Some(spec),
exposure: if tool.defer_loading {
ToolExposure::Deferred
} else {
ToolExposure::Direct
},
search_text: build_dynamic_search_text(tool),
})
}
@@ -64,6 +71,10 @@ impl ToolExecutor<ToolInvocation> for DynamicToolHandler {
self.spec.clone()
}
fn exposure(&self) -> ToolExposure {
self.exposure
}
async fn handle(&self, invocation: ToolInvocation) -> Result<Self::Output, FunctionCallError> {
let ToolInvocation {
session,

View File

@@ -13,6 +13,7 @@ use crate::tools::hook_names::HookToolName;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolExecutor;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolTelemetryTags;
use crate::tools::tool_search_entry::ToolSearchInfo;
@@ -28,11 +29,19 @@ use serde_json::Value;
pub struct McpHandler {
tool_info: ToolInfo,
exposure: ToolExposure,
}
impl McpHandler {
pub fn new(tool_info: ToolInfo) -> Self {
Self { tool_info }
Self::with_exposure(tool_info, ToolExposure::Direct)
}
pub fn with_exposure(tool_info: ToolInfo, exposure: ToolExposure) -> Self {
Self {
tool_info,
exposure,
}
}
}
@@ -71,6 +80,10 @@ impl ToolExecutor<ToolInvocation> for McpHandler {
}))
}
fn exposure(&self) -> ToolExposure {
self.exposure
}
fn supports_parallel_tool_calls(&self) -> bool {
self.tool_info.supports_parallel_tool_calls
}

View File

@@ -35,6 +35,7 @@ use tracing::warn;
pub(crate) type ToolTelemetryTags = Vec<(&'static str, String)>;
pub use codex_tools::ToolExecutor;
pub use codex_tools::ToolExposure;
pub trait ToolHandler: ToolExecutor<ToolInvocation> {
fn search_info(&self) -> Option<ToolSearchInfo> {
@@ -155,11 +156,17 @@ pub(crate) struct PostToolUsePayload {
pub(crate) tool_response: Value,
}
pub(crate) trait AnyToolHandler: Send + Sync {
/// Object-safe registry entry for heterogeneous tool handlers.
///
/// Concrete handlers keep their typed `ToolExecutor::Output`; the registry
/// boxes that output only after typed hooks have run.
pub(crate) trait RegisteredTool: Send + Sync {
fn tool_name(&self) -> ToolName;
fn spec(&self) -> Option<ToolSpec>;
fn exposure(&self) -> ToolExposure;
fn search_info(&self) -> Option<ToolSearchInfo>;
fn supports_parallel_tool_calls(&self) -> bool;
@@ -186,7 +193,7 @@ pub(crate) trait AnyToolHandler: Send + Sync {
) -> BoxFuture<'a, Result<AnyToolResult, FunctionCallError>>;
}
impl<T> AnyToolHandler for T
impl<T> RegisteredTool for T
where
T: ToolHandler,
{
@@ -198,6 +205,10 @@ where
ToolExecutor::spec(self)
}
fn exposure(&self) -> ToolExposure {
ToolExecutor::exposure(self)
}
fn search_info(&self) -> Option<ToolSearchInfo> {
ToolHandler::search_info(self)
}
@@ -253,11 +264,11 @@ where
}
pub struct ToolRegistry {
handlers: HashMap<ToolName, Arc<dyn AnyToolHandler>>,
handlers: HashMap<ToolName, Arc<dyn RegisteredTool>>,
}
impl ToolRegistry {
fn new(handlers: HashMap<ToolName, Arc<dyn AnyToolHandler>>) -> Self {
fn new(handlers: HashMap<ToolName, Arc<dyn RegisteredTool>>) -> Self {
Self { handlers }
}
@@ -272,10 +283,10 @@ impl ToolRegistry {
T: ToolHandler + 'static,
{
let name = handler.tool_name();
Self::new(HashMap::from([(name, handler as Arc<dyn AnyToolHandler>)]))
Self::new(HashMap::from([(name, handler as Arc<dyn RegisteredTool>)]))
}
fn handler(&self, name: &ToolName) -> Option<Arc<dyn AnyToolHandler>> {
fn handler(&self, name: &ToolName) -> Option<Arc<dyn RegisteredTool>> {
self.handlers.get(name).map(Arc::clone)
}
@@ -534,7 +545,7 @@ impl ToolRegistry {
}
pub struct ToolRegistryBuilder {
handlers: HashMap<ToolName, Arc<dyn AnyToolHandler>>,
handlers: HashMap<ToolName, Arc<dyn RegisteredTool>>,
specs: Vec<ToolSpec>,
}
@@ -554,29 +565,28 @@ impl ToolRegistryBuilder {
where
H: ToolHandler + 'static,
{
self.register_any_handler(handler);
self.register_tool(handler);
}
pub(crate) fn register_any_handler(&mut self, handler: Arc<dyn AnyToolHandler>) {
self.register_any_handler_internal(handler, /*include_spec*/ true);
pub(crate) fn register_tool(&mut self, handler: Arc<dyn RegisteredTool>) {
self.register_tool_internal(handler, /*include_spec*/ true);
}
pub(crate) fn register_any_handler_without_spec(&mut self, handler: Arc<dyn AnyToolHandler>) {
self.register_any_handler_internal(handler, /*include_spec*/ false);
pub(crate) fn register_tool_without_spec(&mut self, handler: Arc<dyn RegisteredTool>) {
self.register_tool_internal(handler, /*include_spec*/ false);
}
fn register_any_handler_internal(
&mut self,
handler: Arc<dyn AnyToolHandler>,
include_spec: bool,
) {
fn register_tool_internal(&mut self, handler: Arc<dyn RegisteredTool>, include_spec: bool) {
let name = handler.tool_name();
if self.handlers.contains_key(&name) {
error_or_panic(format!("handler for tool {name} already registered"));
return;
}
if include_spec && let Some(spec) = handler.spec() {
if include_spec
&& handler.exposure() == ToolExposure::Direct
&& let Some(spec) = handler.spec()
{
self.push_spec(spec);
}
@@ -590,12 +600,8 @@ impl ToolRegistryBuilder {
return;
}
if let Some(spec) = executor.spec() {
self.push_spec(spec);
}
let handler: Arc<dyn AnyToolHandler> = Arc::new(ExtensionToolHandler::new(executor));
self.handlers.insert(tool_name, handler);
let handler: Arc<dyn RegisteredTool> = Arc::new(ExtensionToolHandler::new(executor));
self.register_tool_internal(handler, /*include_spec*/ true);
}
pub fn build(self) -> (Vec<ToolSpec>, ToolRegistry) {

View File

@@ -33,10 +33,10 @@ fn handler_looks_up_namespaced_aliases_explicitly() {
let namespaced_name = codex_tools::ToolName::namespaced(namespace, tool_name);
let plain_handler = Arc::new(TestHandler {
tool_name: plain_name.clone(),
}) as Arc<dyn AnyToolHandler>;
}) as Arc<dyn RegisteredTool>;
let namespaced_handler = Arc::new(TestHandler {
tool_name: namespaced_name.clone(),
}) as Arc<dyn AnyToolHandler>;
}) as Arc<dyn RegisteredTool>;
let registry = ToolRegistry::new(HashMap::from([
(plain_name.clone(), Arc::clone(&plain_handler)),
(namespaced_name.clone(), Arc::clone(&namespaced_handler)),

View File

@@ -17,11 +17,9 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::models::SearchToolCallParams;
use codex_protocol::models::ShellToolCallParams;
use codex_tools::DiscoverableTool;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_tools::ToolsConfig;
use std::collections::HashSet;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::instrument;
@@ -66,21 +64,11 @@ impl ToolRouter {
dynamic_tools,
);
let (specs, registry) = builder.build();
let deferred_dynamic_tools = dynamic_tools
.iter()
.filter(|tool| tool.defer_loading)
.map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone()))
.collect::<HashSet<_>>();
let model_visible_specs = specs
.iter()
.filter_map(|spec| {
if config.code_mode_only_enabled
&& codex_code_mode::is_code_mode_nested_tool(spec.name())
{
return None;
}
filter_deferred_dynamic_tool_spec(spec.clone(), &deferred_dynamic_tools)
.into_iter()
.filter(|spec| {
!config.code_mode_only_enabled
|| !codex_code_mode::is_code_mode_nested_tool(spec.name())
})
.collect();
@@ -232,38 +220,6 @@ pub(crate) fn extension_tool_executors(session: &Session) -> Vec<Arc<dyn Extensi
.collect()
}
fn filter_deferred_dynamic_tool_spec(
spec: ToolSpec,
deferred_dynamic_tools: &HashSet<ToolName>,
) -> Option<ToolSpec> {
if deferred_dynamic_tools.is_empty() {
return Some(spec);
}
match spec {
ToolSpec::Function(tool) => {
if deferred_dynamic_tools.contains(&ToolName::plain(tool.name.as_str())) {
None
} else {
Some(ToolSpec::Function(tool))
}
}
ToolSpec::Namespace(mut namespace) => {
let namespace_name = namespace.name.clone();
namespace.tools.retain(|tool| match tool {
ResponsesApiNamespaceTool::Function(tool) => !deferred_dynamic_tools.contains(
&ToolName::namespaced(namespace_name.as_str(), tool.name.as_str()),
),
});
if namespace.tools.is_empty() {
None
} else {
Some(ToolSpec::Namespace(namespace))
}
}
spec => Some(spec),
}
}
#[cfg(test)]
#[path = "router_tests.rs"]
mod tests;

View File

@@ -44,7 +44,8 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions;
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::AnyToolHandler;
use crate::tools::registry::RegisteredTool;
use crate::tools::registry::ToolExposure;
use crate::tools::registry::ToolRegistryBuilder;
use crate::tools::spec_plan_types::ToolRegistryBuildParams;
use crate::tools::spec_plan_types::agent_type_description;
@@ -52,13 +53,11 @@ use codex_extension_api::ExtensionToolExecutor;
use codex_protocol::openai_models::ConfigShellToolType;
use codex_tools::ResponsesApiNamespaceTool;
use codex_tools::ToolEnvironmentMode;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_tools::ToolsConfig;
use codex_tools::collect_code_mode_exec_prompt_tool_definitions;
use codex_tools::default_namespace_description;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::sync::Arc;
pub fn build_tool_registry_builder(
@@ -66,40 +65,34 @@ pub fn build_tool_registry_builder(
params: ToolRegistryBuildParams<'_>,
) -> ToolRegistryBuilder {
let mut builder = ToolRegistryBuilder::new();
let all_deferred_tools = params
.deferred_mcp_tools
.into_iter()
.flatten()
.map(codex_mcp::ToolInfo::canonical_tool_name)
.chain(
params
.dynamic_tools
.iter()
.filter(|tool| tool.defer_loading)
.map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone())),
)
.collect::<HashSet<_>>();
let handlers = collect_handler_tools(config, params);
let deferred_tools_available = handlers
.iter()
.any(|handler| handler.exposure() == ToolExposure::Deferred);
for handler in build_code_mode_handlers(
config,
&handlers,
params.extension_tool_executors,
config.search_tool && !all_deferred_tools.is_empty(),
config.search_tool && deferred_tools_available,
) {
builder.register_any_handler(handler);
builder.register_tool(handler);
}
let mut non_deferred_specs = Vec::new();
let mut deferred_search_infos = Vec::new();
for handler in &handlers {
let tool_name = handler.tool_name();
if all_deferred_tools.contains(&tool_name) {
if let Some(search_info) = handler.search_info() {
deferred_search_infos.push(search_info);
match handler.exposure() {
ToolExposure::Direct => {
if let Some(spec) = handler.spec() {
non_deferred_specs.push(spec);
}
}
ToolExposure::Deferred => {
if let Some(search_info) = handler.search_info() {
deferred_search_infos.push(search_info);
}
}
} else if let Some(spec) = handler.spec() {
non_deferred_specs.push(spec);
}
}
@@ -127,7 +120,7 @@ pub fn build_tool_registry_builder(
}
for handler in handlers {
builder.register_any_handler_without_spec(handler);
builder.register_tool_without_spec(handler);
}
if config.search_tool && config.namespace_tools && !deferred_search_infos.is_empty() {
@@ -143,10 +136,10 @@ pub fn build_tool_registry_builder(
fn build_code_mode_handlers(
config: &ToolsConfig,
handlers: &[Arc<dyn AnyToolHandler>],
handlers: &[Arc<dyn RegisteredTool>],
extension_tool_executors: &[Arc<dyn ExtensionToolExecutor>],
deferred_tools_available: bool,
) -> Vec<Arc<dyn AnyToolHandler>> {
) -> Vec<Arc<dyn RegisteredTool>> {
if !config.code_mode_enabled {
return vec![];
}
@@ -251,9 +244,9 @@ fn code_mode_namespace_descriptions(
fn collect_handler_tools(
config: &ToolsConfig,
params: ToolRegistryBuildParams<'_>,
) -> Vec<Arc<dyn AnyToolHandler>> {
) -> Vec<Arc<dyn RegisteredTool>> {
let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled;
let mut handlers = Vec::<Arc<dyn AnyToolHandler>>::new();
let mut handlers = Vec::<Arc<dyn RegisteredTool>>::new();
if config.environment_mode.has_environment() {
let include_environment_id =
@@ -430,7 +423,10 @@ fn collect_handler_tools(
if let Some(deferred_mcp_tools) = params.deferred_mcp_tools {
for tool in deferred_mcp_tools {
handlers.push(Arc::new(McpHandler::new(tool.clone())));
handlers.push(Arc::new(McpHandler::with_exposure(
tool.clone(),
ToolExposure::Deferred,
)));
}
}

View File

@@ -79,6 +79,7 @@ pub use tool_discovery::ToolSearchSourceInfo;
pub use tool_discovery::collect_request_plugin_install_entries;
pub use tool_discovery::filter_request_plugin_install_discoverable_tools_for_client;
pub use tool_executor::ToolExecutor;
pub use tool_executor::ToolExposure;
pub use tool_output::JsonToolOutput;
pub use tool_output::ToolOutput;
pub use tool_payload::ToolPayload;

View File

@@ -5,6 +5,14 @@ use crate::ToolName;
use crate::ToolOutput;
use crate::ToolSpec;
/// Controls whether a tool is exposed in the initial model-visible tool list
/// or registered for later discovery.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ToolExposure {
Direct,
Deferred,
}
/// Shared runtime contract for model-visible tools.
///
/// Implementations keep the model-visible spec tied to the executable runtime.
@@ -20,6 +28,10 @@ pub trait ToolExecutor<Invocation>: Send + Sync {
None
}
fn exposure(&self) -> ToolExposure {
ToolExposure::Direct
}
fn supports_parallel_tool_calls(&self) -> bool {
false
}