From b293412c2434f4351c1c50a946fb0d4964ce0eef Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 30 Jul 2026 09:14:05 +0000 Subject: [PATCH] Unify tool runtime assembly before router construction (#36119) ## What changed - Split core tool planning from the assembly of MCP, extension, and dynamic tool runtimes. - Build `ToolRouter` from a single ordered runtime list plus hosted tool specs. - Centralize whether each `ToolExposure` variant is available in code mode. - Preserve core, MCP, extension, and dynamic source ordering and existing name-collision priority. ## Testing - Add coverage for unified source ordering and collisions, code-mode exposure, and standalone web search alongside MCP and dynamic `web.run` tools. GitOrigin-RevId: b010c2e9728b92b75a35325ca9b2ec465c3bac12 --- codex-rs/core/src/session/tests.rs | 50 ++-- codex-rs/core/src/session/turn.rs | 35 ++- .../core/src/stream_events_utils_tests.rs | 15 +- codex-rs/core/src/tools/router.rs | 44 +--- codex-rs/core/src/tools/router_tests.rs | 150 ++++++----- codex-rs/core/src/tools/spec_plan.rs | 212 ++++++++-------- codex-rs/core/src/tools/spec_plan_tests.rs | 233 +++++++++++++----- codex-rs/tools/src/tool_executor.rs | 8 + 8 files changed, 415 insertions(+), 332 deletions(-) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 1ec3e73b6e..b33c799224 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -18,6 +18,7 @@ use crate::skills::SkillRenderSideEffects; use crate::skills::render::SkillMetadataBudget; 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; @@ -74,6 +75,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; +use codex_tools::ToolSpec; use codex_utils_path_uri::PathUri; use tracing::Span; @@ -815,19 +817,33 @@ async fn preview_session_start_hooks( ) } -fn test_tool_runtime(session: Arc, turn_context: Arc) -> ToolCallRuntime { - let step_context = StepContext::for_test(Arc::clone(&turn_context)); - let router = Arc::new(ToolRouter::from_context( +pub(crate) fn tool_runtimes_for_test_step( + step_context: &StepContext, +) -> (Vec>, Vec) { + let mut tool_runtimes = crate::tools::spec_plan::build_core_tool_runtimes( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), - crate::tools::router::ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn_context.dynamic_tools.as_slice(), - }, + /*tool_suggest_candidates*/ None, + /*wait_for_environment_tool_config*/ None, + ); + let hosted_specs = crate::tools::spec_plan::append_source_tool_runtimes( + step_context.turn.as_ref(), + &mut tool_runtimes, + Vec::new(), + Vec::new(), + &step_context.turn.dynamic_tools, + ); + (tool_runtimes, hosted_specs) +} + +fn test_tool_runtime(session: Arc, turn_context: Arc) -> 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( + step_context.turn.as_ref(), + tool_runtimes, + hosted_specs, &Default::default(), )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); @@ -11087,17 +11103,11 @@ 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 router = ToolRouter::from_context( + let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref()); + let router = ToolRouter::from_tools( step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - crate::tools::router::ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn_context.dynamic_tools.as_slice(), - }, + tool_runtimes, + hosted_specs, &Default::default(), ); let item = ResponseItem::CustomToolCall { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 57691ab81a..947de8226f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -58,10 +58,11 @@ use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; use crate::tools::registry::ToolArgumentDiffConsumer; -use crate::tools::router::ToolRouterParams; use crate::tools::router::ToolSuggestCandidates; use crate::tools::router::ToolSuggestPresentation; -use crate::tools::router::extension_tool_executors; +use crate::tools::spec_plan::append_source_tool_runtimes; +use crate::tools::spec_plan::build_core_tool_runtimes; +use crate::tools::spec_plan::extension_tool_executors; use crate::tools::spec_plan::search_tool_enabled; use crate::tools::spec_plan::tool_suggest_enabled; use crate::turn_diff_tracker::TurnDiffTracker; @@ -1554,20 +1555,28 @@ pub(crate) async fn built_tools( &turn_context.config, search_tool_enabled(turn_context), ); - let tool_router = Arc::new(ToolRouter::from_context( + let wait_for_environment_tool_config = sess + .services + .thread_extension_data + .get::(); + let mut tool_runtimes = build_core_tool_runtimes( turn_context, environments, mcp, - ToolRouterParams { - tool_runtimes: mcp_tool_runtimes, - tool_suggest_candidates, - extension_tool_executors: extension_tool_executors(sess, step_store), - wait_for_environment_tool_config: sess - .services - .thread_extension_data - .get::(), - dynamic_tools: turn_context.dynamic_tools.as_slice(), - }, + tool_suggest_candidates.as_ref(), + wait_for_environment_tool_config.as_ref(), + ); + let hosted_specs = append_source_tool_runtimes( + turn_context, + &mut tool_runtimes, + mcp_tool_runtimes, + extension_tool_executors(sess, step_store), + &turn_context.dynamic_tools, + ); + let tool_router = Arc::new(ToolRouter::from_tools( + turn_context, + tool_runtimes, + hosted_specs, &sess.services.tool_search_handler_cache, )); (all_mcp_tools, tool_router) diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 9485824adb..94e370009c 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -8,6 +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::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; use crate::turn_diff_tracker::TurnDiffTracker; @@ -277,17 +278,11 @@ 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 router = Arc::new(ToolRouter::from_context( + let (tool_runtimes, hosted_specs) = tool_runtimes_for_test_step(step_context.as_ref()); + let router = Arc::new(ToolRouter::from_tools( step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - crate::tools::router::ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn_context.dynamic_tools.as_slice(), - }, + tool_runtimes, + hosted_specs, &Default::default(), )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 73c98dec11..1e604297af 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,4 +1,3 @@ -use crate::environment_selection::TurnEnvironmentSnapshot; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::step_context::StepContext; @@ -12,12 +11,9 @@ use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolRegistry; use crate::tools::spec_plan::build_tool_router; -use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::ResponseItem; use codex_protocol::models::SearchToolCallParams; use codex_tools::DiscoverableTool; -use codex_tools::ToolCall as ExtensionToolCall; -use codex_tools::ToolExecutor; use codex_tools::ToolName; use codex_tools::ToolSpec; use std::borrow::Cow; @@ -71,14 +67,6 @@ pub struct ToolRouter { model_visible_specs: Vec, } -pub(crate) struct ToolRouterParams<'a> { - pub(crate) tool_runtimes: Vec>, - pub(crate) tool_suggest_candidates: Option, - pub(crate) extension_tool_executors: Vec>>, - pub(crate) wait_for_environment_tool_config: Option>, - pub(crate) dynamic_tools: &'a [DynamicToolSpec], -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ToolSuggestPresentation { ListTool, @@ -92,18 +80,16 @@ pub(crate) struct ToolSuggestCandidates { } impl ToolRouter { - pub(crate) fn from_context( + pub(crate) fn from_tools( turn_context: &TurnContext, - environments: &TurnEnvironmentSnapshot, - mcp: &codex_mcp::McpBinding, - params: ToolRouterParams<'_>, + tool_runtimes: Vec>, + hosted_specs: Vec, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { build_tool_router( turn_context, - environments, - mcp, - params, + tool_runtimes, + hosted_specs, tool_search_handler_cache, ) } @@ -299,26 +285,6 @@ impl ToolRouter { } } -#[instrument(level = "trace", skip_all)] -pub(crate) fn extension_tool_executors( - session: &Session, - step_store: &codex_extension_api::ExtensionData, -) -> Vec>> { - session - .services - .extensions - .tool_contributors() - .iter() - .flat_map(|contributor| { - contributor.tools_for_step( - &session.services.session_extension_data, - &session.services.thread_extension_data, - step_store, - ) - }) - .collect() -} - #[cfg(test)] #[path = "router_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 0e4dc01cf7..3abe1cdb41 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -9,6 +9,9 @@ use crate::tools::handlers::McpHandler; use crate::tools::registry::CoreToolRuntime; 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::extension_tool_executors; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistry; @@ -37,8 +40,6 @@ use tokio_util::sync::CancellationToken; use super::ToolCall; use super::ToolCallSource; use super::ToolRouter; -use super::ToolRouterParams; -use super::extension_tool_executors; use super::tool_log_payload; struct ExtensionEchoContributor; @@ -125,23 +126,44 @@ fn extension_tool_test_registry() -> Arc> { Arc::new(builder.build()) } +fn test_tool_router( + step_context: &StepContext, + mcp_tool_runtimes: Vec>, + extension_tool_executors: Vec>>, + dynamic_tools: &[DynamicToolSpec], +) -> ToolRouter { + let mut tool_runtimes = build_core_tool_runtimes( + 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( + step_context.turn.as_ref(), + &mut tool_runtimes, + mcp_tool_runtimes, + extension_tool_executors, + dynamic_tools, + ); + ToolRouter::from_tools( + step_context.turn.as_ref(), + tool_runtimes, + hosted_specs, + &Default::default(), + ) +} + #[tokio::test] async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let router = ToolRouter::from_context( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn.dynamic_tools.as_slice(), - }, - &Default::default(), + let router = test_tool_router( + step_context.as_ref(), + Vec::new(), + Vec::new(), + &turn.dynamic_tools, ); let parallel_tool_name = ["exec_command", "shell_command"] @@ -249,34 +271,27 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let router = ToolRouter::from_context( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: vec![ + let router = test_tool_router( + step_context.as_ref(), + vec![ + mcp_runtime(mcp_tool_info( + "echo", + /*supports_parallel_tool_calls*/ true, + "mcp__echo__", + "query_with_delay", + )), + override_tool_exposure( mcp_runtime(mcp_tool_info( - "echo", - /*supports_parallel_tool_calls*/ true, - "mcp__echo__", + "hello_echo", + /*supports_parallel_tool_calls*/ false, + "mcp__hello_echo__", "query_with_delay", )), - override_tool_exposure( - mcp_runtime(mcp_tool_info( - "hello_echo", - /*supports_parallel_tool_calls*/ false, - "mcp__hello_echo__", - "query_with_delay", - )), - ToolExposure::DirectModelOnly, - ), - ], - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn.dynamic_tools.as_slice(), - }, - &Default::default(), + ToolExposure::DirectModelOnly, + ), + ], + Vec::new(), + &turn.dynamic_tools, ); let call = ToolCall { @@ -312,18 +327,11 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> let (_, turn) = make_session_and_context().await; let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let router = ToolRouter::from_context( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: turn.dynamic_tools.as_slice(), - }, - &Default::default(), + let router = test_tool_router( + step_context.as_ref(), + Vec::new(), + Vec::new(), + &turn.dynamic_tools, ); assert!(!router.tool_supports_parallel(&ToolCall { @@ -372,18 +380,11 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { ], })]; - let router = ToolRouter::from_context( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: &dynamic_tools, - }, - &Default::default(), + let router = test_tool_router( + step_context.as_ref(), + Vec::new(), + Vec::new(), + &dynamic_tools, ); assert_eq!( @@ -450,21 +451,14 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow let mut expected_history_item = history_item.clone(); expected_history_item.set_turn_id_if_missing(&turn.sub_id); - let router = ToolRouter::from_context( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: None, - tool_runtimes: Vec::new(), - extension_tool_executors: extension_tool_executors( - &session, - &codex_extension_api::ExtensionData::new(turn.sub_id.clone()), - ), - wait_for_environment_tool_config: None, - dynamic_tools: turn.dynamic_tools.as_slice(), - }, - &Default::default(), + let router = test_tool_router( + step_context.as_ref(), + Vec::new(), + extension_tool_executors( + &session, + &codex_extension_api::ExtensionData::new(turn.sub_id.clone()), + ), + &turn.dynamic_tools, ); assert!( diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 5ff9e392f3..215ea66463 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1,6 +1,7 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::default_exec_yield_time_override_ms; use crate::tools::code_mode::execute_spec::create_code_mode_tool; @@ -57,7 +58,7 @@ use crate::tools::registry::ToolExposure; use crate::tools::registry::ToolRegistry; use crate::tools::registry::override_tool_exposure; use crate::tools::router::ToolRouter; -use crate::tools::router::ToolRouterParams; +use codex_extension_api::ExtensionData; use codex_features::Feature; use codex_login::AuthManager; use codex_protocol::account::PlanType; @@ -130,10 +131,6 @@ impl PlannedTools { self.add_with_exposure(handler, ToolExposure::Hidden); } - fn add_hosted_spec(&mut self, spec: ToolSpec) { - self.hosted_specs.push(spec); - } - fn runtimes(&self) -> &[PlannedRuntime] { &self.runtimes } @@ -144,70 +141,92 @@ struct CoreToolPlanContext<'a> { turn_context: &'a TurnContext, environments: &'a TurnEnvironmentSnapshot, mcp: &'a codex_mcp::McpBinding, - tool_runtimes: &'a [PlannedRuntime], tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, - extension_tool_executors: &'a [Arc>], wait_for_environment_tool_config: Option<&'a Arc>, - dynamic_tools: &'a [DynamicToolSpec], - tool_search_handler_cache: &'a ToolSearchHandlerCache, default_agent_type_description: &'a str, wait_agent_timeouts: WaitAgentTimeoutOptions, } #[instrument(level = "trace", skip_all)] -pub(crate) fn build_tool_router( +pub(crate) fn build_core_tool_runtimes( turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, mcp: &codex_mcp::McpBinding, - params: ToolRouterParams<'_>, - tool_search_handler_cache: &ToolSearchHandlerCache, -) -> ToolRouter { - let (model_visible_specs, registry) = build_tool_specs_and_registry( - turn_context, - environments, - mcp, - params, - tool_search_handler_cache, - ); - ToolRouter::from_parts(registry, model_visible_specs) -} - -#[instrument(level = "trace", skip_all)] -fn build_tool_specs_and_registry( - turn_context: &TurnContext, - environments: &TurnEnvironmentSnapshot, - mcp: &codex_mcp::McpBinding, - params: ToolRouterParams<'_>, - tool_search_handler_cache: &ToolSearchHandlerCache, -) -> (Vec, ToolRegistry) { - let ToolRouterParams { - tool_runtimes, - tool_suggest_candidates, - extension_tool_executors, - wait_for_environment_tool_config, - dynamic_tools, - } = params; + tool_suggest_candidates: Option<&crate::tools::router::ToolSuggestCandidates>, + wait_for_environment_tool_config: Option<&Arc>, +) -> Vec> { let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { turn_context, environments, mcp, - tool_runtimes: &tool_runtimes, - tool_suggest_candidates: tool_suggest_candidates.as_ref(), - extension_tool_executors: &extension_tool_executors, - wait_for_environment_tool_config: wait_for_environment_tool_config.as_ref(), - dynamic_tools, - tool_search_handler_cache, + tool_suggest_candidates, + wait_for_environment_tool_config, default_agent_type_description: &default_agent_type_description, wait_agent_timeouts: wait_agent_timeout_options(turn_context), }; let mut planned_tools = PlannedTools::default(); - add_tool_sources(&context, &mut planned_tools); + add_core_tool_sources(&context, &mut planned_tools); + planned_tools.runtimes +} + +#[instrument(level = "trace", skip_all)] +pub(crate) fn append_source_tool_runtimes( + turn_context: &TurnContext, + tool_runtimes: &mut Vec>, + mcp_tool_runtimes: Vec>, + extension_tool_executors: Vec>>, + dynamic_tools: &[DynamicToolSpec], +) -> Vec { + if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { + return Vec::new(); + } + + let hosted_specs = hosted_model_tool_specs(turn_context, &extension_tool_executors); + tool_runtimes.extend(mcp_tool_runtimes); + append_extension_tool_executors(turn_context, &extension_tool_executors, tool_runtimes); + append_dynamic_tool_runtimes(dynamic_tools, tool_runtimes); + hosted_specs +} + +#[instrument(level = "trace", skip_all)] +pub(crate) fn extension_tool_executors( + session: &Session, + step_store: &ExtensionData, +) -> Vec>> { + session + .services + .extensions + .tool_contributors() + .iter() + .flat_map(|contributor| { + contributor.tools_for_step( + &session.services.session_extension_data, + &session.services.thread_extension_data, + step_store, + ) + }) + .collect() +} + +#[instrument(level = "trace", skip_all)] +pub(crate) fn build_tool_router( + turn_context: &TurnContext, + tool_runtimes: Vec>, + hosted_specs: Vec, + tool_search_handler_cache: &ToolSearchHandlerCache, +) -> ToolRouter { + let mut planned_tools = PlannedTools { + runtimes: tool_runtimes, + hosted_specs, + }; apply_direct_model_only_namespace_overrides(turn_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) + append_tool_search_executor(turn_context, &mut planned_tools, tool_search_handler_cache); + prepend_code_mode_executors(turn_context, &mut planned_tools); + let (model_visible_specs, registry) = + build_model_visible_specs_and_registry(turn_context, planned_tools); + ToolRouter::from_parts(registry, model_visible_specs) } fn apply_direct_model_only_namespace_overrides( @@ -226,15 +245,8 @@ fn apply_direct_model_only_namespace_overrides( .direct_only_tool_namespaces .contains(namespace) }); - match runtime.exposure() { - ToolExposure::Direct | ToolExposure::Deferred if configured => { - *runtime = - override_tool_exposure(Arc::clone(runtime), ToolExposure::DirectModelOnly); - } - ToolExposure::Direct - | ToolExposure::Deferred - | ToolExposure::DirectModelOnly - | ToolExposure::Hidden => {} + if configured && runtime.exposure().is_available_in_code_mode() { + *runtime = override_tool_exposure(Arc::clone(runtime), ToolExposure::DirectModelOnly); } } } @@ -288,7 +300,7 @@ fn spec_for_model_request( ) -> ToolSpec { let tool_mode = effective_tool_mode(turn_context); if matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - && exposure != ToolExposure::DirectModelOnly + && exposure.is_available_in_code_mode() && !is_excluded_from_code_mode(turn_context, tool_name) && codex_code_mode::is_code_mode_nested_tool(spec.name()) { @@ -299,17 +311,20 @@ fn spec_for_model_request( } #[instrument(level = "trace", skip_all)] -fn hosted_model_tool_specs(context: &CoreToolPlanContext<'_>) -> Vec { - let turn_context = context.turn_context; +fn hosted_model_tool_specs( + turn_context: &TurnContext, + extension_tool_executors: &[Arc>], +) -> Vec { // Responses Lite accepts schemas for client-executed tools, not hosted Responses tools. - if turn_context.model_info.use_responses_lite { + if turn_context.model_info.use_responses_lite + || crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) + { return Vec::new(); } let mut specs = Vec::new(); let standalone_web_search_available = standalone_web_search_enabled(turn_context) - && context - .extension_tool_executors + && extension_tool_executors .iter() .any(|executor| executor.tool_name() == ToolName::namespaced("web", "run")); // `Some(Cached/Live/Disabled)` are the options for mode when standalone search is unavailable @@ -438,7 +453,7 @@ fn is_hidden_by_code_mode_only( ) -> bool { let tool_mode = effective_tool_mode(turn_context); tool_mode == ToolMode::CodeModeOnly - && exposure != ToolExposure::DirectModelOnly + && exposure.is_available_in_code_mode() && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name( tool_name, )) @@ -469,11 +484,7 @@ fn build_code_mode_executors( let deferred_tools_guidance_enabled = search_tool_enabled(turn_context); for executor in executors { let exposure = executor.exposure(); - if exposure == ToolExposure::DirectModelOnly { - continue; - } - - if exposure == ToolExposure::Hidden { + if !exposure.is_available_in_code_mode() { continue; } @@ -605,7 +616,7 @@ fn code_mode_namespace_descriptions( } #[instrument(level = "trace", skip_all)] -fn add_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { +fn add_core_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { // Guardian reviewers receive only `exec_command`, `write_stdin`, and `view_image` // when an environment is available; all general tool sources stay excluded. if crate::guardian::is_guardian_reviewer_source(&context.turn_context.session_source) { @@ -637,14 +648,6 @@ fn add_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut Plann add_mcp_resource_tools(context, planned_tools); add_core_utility_tools(context, planned_tools); add_collaboration_tools(context, planned_tools); - for runtime in context.tool_runtimes { - planned_tools.add_arc(Arc::clone(runtime)); - } - add_extension_tools(context, planned_tools); - add_dynamic_tools(context, planned_tools); - for spec in hosted_model_tool_specs(context) { - planned_tools.add_hosted_spec(spec); - } } fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool { @@ -910,9 +913,12 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu } } -#[instrument(level = "trace", skip_all, fields(dynamic_tool_count = context.dynamic_tools.len()))] -fn add_dynamic_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { - for spec in context.dynamic_tools { +#[instrument(level = "trace", skip_all, fields(dynamic_tool_count = dynamic_tools.len()))] +fn append_dynamic_tool_runtimes( + dynamic_tools: &[DynamicToolSpec], + tool_runtimes: &mut Vec>, +) { + for spec in dynamic_tools { match spec { DynamicToolSpec::Function(tool) => { let Some(handler) = DynamicToolHandler::new(tool) else { @@ -922,7 +928,7 @@ fn add_dynamic_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Plan ); continue; }; - planned_tools.add(handler); + tool_runtimes.push(Arc::new(handler)); } DynamicToolSpec::Namespace(namespace) => { for tool in &namespace.tools { @@ -936,34 +942,19 @@ fn add_dynamic_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Plan ); continue; }; - planned_tools.add(handler); + tool_runtimes.push(Arc::new(handler)); } } } } } -#[instrument( - level = "trace", - skip_all, - fields(extension_tool_executor_count = context.extension_tool_executors.len()) -)] -fn add_extension_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { - // Extension ToolContributor implementations are resolved into executors - // before planning. Core only adapts those executors into its runtime set. - append_extension_tool_executors( - context.turn_context, - context.extension_tool_executors, - planned_tools, - ); -} - #[instrument(level = "trace", skip_all)] fn append_tool_search_executor( - context: &CoreToolPlanContext<'_>, + turn_context: &TurnContext, planned_tools: &mut PlannedTools, + tool_search_handler_cache: &ToolSearchHandlerCache, ) { - let turn_context = context.turn_context; if !search_tool_enabled(turn_context) { return; } @@ -987,17 +978,12 @@ fn append_tool_search_executor( } else { ToolSearchSourceListing::Include }; - let handler: PlannedRuntime = context - .tool_search_handler_cache - .get_or_build(search_infos, source_listing); + let handler: PlannedRuntime = + tool_search_handler_cache.get_or_build(search_infos, source_listing); planned_tools.add_arc(handler); } -fn prepend_code_mode_executors( - context: &CoreToolPlanContext<'_>, - planned_tools: &mut PlannedTools, -) { - let turn_context = context.turn_context; +fn prepend_code_mode_executors(turn_context: &TurnContext, planned_tools: &mut PlannedTools) { let code_mode_executors = build_code_mode_executors(turn_context, planned_tools.runtimes()); planned_tools.runtimes.splice(0..0, code_mode_executors); } @@ -1005,14 +991,13 @@ fn prepend_code_mode_executors( fn append_extension_tool_executors( turn_context: &TurnContext, executors: &[Arc>], - planned_tools: &mut PlannedTools, + tool_runtimes: &mut Vec>, ) { if executors.is_empty() { return; } - let mut reserved_tool_names = planned_tools - .runtimes() + let mut reserved_tool_names = tool_runtimes .iter() .map(|executor| executor.tool_name()) .collect::>(); @@ -1022,8 +1007,7 @@ fn append_extension_tool_executors( reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME)); } if search_tool_enabled(turn_context) - && planned_tools - .runtimes() + && tool_runtimes .iter() .any(|executor| executor.exposure() == ToolExposure::Deferred) { @@ -1049,7 +1033,7 @@ fn append_extension_tool_executors( warn!("Skipping extension tool `{tool_name}`: tool already registered"); continue; } - planned_tools.add(ExtensionToolAdapter::new(executor)); + tool_runtimes.push(Arc::new(ExtensionToolAdapter::new(executor))); } } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 87be4d9275..e1b3ab1ca9 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -42,9 +42,10 @@ 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::router::ToolRouter; -use crate::tools::router::ToolRouterParams; 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; const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration"; @@ -188,17 +189,24 @@ async fn probe_with( configure_turn(&mut turn); let turn = Arc::new(turn); let step_context = StepContext::for_test(Arc::clone(&turn)); - let router = ToolRouter::from_context( + let mut tool_runtimes = build_core_tool_runtimes( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), - ToolRouterParams { - tool_suggest_candidates: inputs.tool_suggest_candidates, - tool_runtimes: inputs.tool_runtimes, - extension_tool_executors: inputs.extension_tool_executors, - wait_for_environment_tool_config: inputs.wait_for_environment_tool_config, - dynamic_tools: inputs.dynamic_tools.as_slice(), - }, + inputs.tool_suggest_candidates.as_ref(), + inputs.wait_for_environment_tool_config.as_ref(), + ); + let hosted_specs = append_source_tool_runtimes( + step_context.turn.as_ref(), + &mut tool_runtimes, + inputs.tool_runtimes, + inputs.extension_tool_executors, + &inputs.dynamic_tools, + ); + let router = ToolRouter::from_tools( + step_context.turn.as_ref(), + tool_runtimes, + hosted_specs, &Default::default(), ); ToolPlanProbe::from_router(router) @@ -829,17 +837,16 @@ async fn environment_tools_follow_the_step_context() { &turn.config, ))); - let plan = ToolPlanProbe::from_router(ToolRouter::from_context( + let plan = ToolPlanProbe::from_router(ToolRouter::from_tools( turn.as_ref(), - &environments, - mcp.as_ref(), - ToolRouterParams { - tool_runtimes: Vec::new(), - tool_suggest_candidates: None, - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: &[], - }, + build_core_tool_runtimes( + turn.as_ref(), + &environments, + mcp.as_ref(), + /*tool_suggest_candidates*/ None, + /*wait_for_environment_tool_config*/ None, + ), + super::hosted_model_tool_specs(turn.as_ref(), &[]), &Default::default(), )); @@ -1000,6 +1007,78 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { ]); } +#[tokio::test] +async fn unified_tool_runtimes_preserve_source_order_and_collision_priority() { + let plan = probe_with( + |turn| { + set_features(turn, &[Feature::ShellTool, Feature::UnifiedExec]); + set_feature(turn, Feature::ShellZshFork, /*enabled*/ false); + turn.model_info.shell_type = ConfigShellToolType::ShellCommand; + }, + ToolPlanInputs { + tool_runtimes: vec![mcp_runtime( + "registry", + "mcp__registry", + "lookup", + ToolExposure::Direct, + )], + extension_tool_executors: vec![ + Arc::new(TestNamespaceExtensionTool { + namespace: "mcp__registry", + tool_name: "lookup", + }), + Arc::new(TestNamespaceExtensionTool { + namespace: "registry_extension", + tool_name: "lookup", + }), + ], + dynamic_tools: vec![dynamic_tool( + Some("registry_dynamic"), + "lookup", + /*defer_loading*/ false, + )], + ..ToolPlanInputs::default() + }, + ) + .await; + + let expected_source_order = [ + "exec_command", + "mcp__registry", + "registry_extension", + "registry_dynamic", + ]; + let source_order = plan + .visible_names + .iter() + .map(String::as_str) + .filter(|name| expected_source_order.contains(name)) + .collect::>(); + assert_eq!(source_order, expected_source_order); + + let mcp_tool_name = ToolName::namespaced("mcp__registry", "lookup").to_string(); + let extension_tool_name = ToolName::namespaced("registry_extension", "lookup").to_string(); + let dynamic_tool_name = ToolName::namespaced("registry_dynamic", "lookup").to_string(); + plan.assert_registered_contains(&[ + "exec_command", + &mcp_tool_name, + &extension_tool_name, + &dynamic_tool_name, + ]); + assert_eq!( + plan.namespace_function_names("mcp__registry"), + &["lookup".to_string()] + ); + + let ToolSpec::Namespace(namespace) = plan.visible_spec("mcp__registry") else { + panic!("expected the MCP namespace to stay visible"); + }; + let [ResponsesApiNamespaceTool::Function(tool)] = namespace.tools.as_slice() else { + panic!("expected exactly one MCP tool after the extension collision"); + }; + assert_eq!(tool.description, "lookup test tool"); +} + #[tokio::test] async fn deferred_extension_tools_are_discoverable_with_tool_search() { let plan = probe_with( @@ -1027,22 +1106,23 @@ 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 first_router = ToolRouter::from_context( + let mut first_tool_runtimes = build_core_tool_runtimes( first_step_context.turn.as_ref(), &first_step_context.environments, first_step_context.mcp.as_ref(), - ToolRouterParams { - tool_runtimes: vec![mcp_runtime( - "first", - "mcp__first", - "lookup", - ToolExposure::Deferred, - )], - tool_suggest_candidates: None, - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: &[], - }, + /*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( + first_step_context.turn.as_ref(), + first_tool_runtimes, + super::hosted_model_tool_specs(first_step_context.turn.as_ref(), &[]), &cache, ); let first_plan = ToolPlanProbe::from_router(first_router); @@ -1051,22 +1131,23 @@ 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 second_router = ToolRouter::from_context( + let mut second_tool_runtimes = build_core_tool_runtimes( second_step_context.turn.as_ref(), &second_step_context.environments, second_step_context.mcp.as_ref(), - ToolRouterParams { - tool_runtimes: vec![mcp_runtime( - "second", - "mcp__second", - "lookup", - ToolExposure::Deferred, - )], - tool_suggest_candidates: None, - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: &[], - }, + /*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( + second_step_context.turn.as_ref(), + second_tool_runtimes, + super::hosted_model_tool_specs(second_step_context.turn.as_ref(), &[]), &cache, ); let second_plan = ToolPlanProbe::from_router(second_router); @@ -1106,22 +1187,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 router = ToolRouter::from_context( + let mut tool_runtimes = build_core_tool_runtimes( step_context.turn.as_ref(), &step_context.environments, step_context.mcp.as_ref(), - ToolRouterParams { - tool_runtimes: vec![mcp_runtime( - "calendar", - "mcp__calendar", - "lookup", - ToolExposure::Deferred, - )], - tool_suggest_candidates: None, - extension_tool_executors: Vec::new(), - wait_for_environment_tool_config: None, - dynamic_tools: &[], - }, + /*tool_suggest_candidates*/ None, + /*wait_for_environment_tool_config*/ None, + ); + tool_runtimes.push(mcp_runtime( + "calendar", + "mcp__calendar", + "lookup", + ToolExposure::Deferred, + )); + let router = ToolRouter::from_tools( + step_context.turn.as_ref(), + tool_runtimes, + super::hosted_model_tool_specs(step_context.turn.as_ref(), &[]), &cache, ); let plan = ToolPlanProbe::from_router(router); @@ -1902,6 +1984,41 @@ async fn hosted_web_search_and_standalone_image_generation_follow_runtime_gates( .await; standalone_web_search_without_web_run.assert_visible_contains(&["web_search"]); + let standalone_web_search_with_dynamic_web_run = probe_with( + |turn| { + set_feature(turn, Feature::StandaloneWebSearch, /*enabled*/ true); + set_web_search_mode(turn, WebSearchMode::Live); + }, + ToolPlanInputs { + dynamic_tools: vec![dynamic_tool( + Some("web"), + "run", + /*defer_loading*/ false, + )], + ..Default::default() + }, + ) + .await; + standalone_web_search_with_dynamic_web_run.assert_visible_contains(&["web", "web_search"]); + + let standalone_web_search_with_mcp_web_run = probe_with( + |turn| { + set_feature(turn, Feature::StandaloneWebSearch, /*enabled*/ true); + set_web_search_mode(turn, WebSearchMode::Live); + }, + ToolPlanInputs { + tool_runtimes: vec![mcp_runtime( + "web_server", + "web", + "run", + ToolExposure::Direct, + )], + ..Default::default() + }, + ) + .await; + standalone_web_search_with_mcp_web_run.assert_visible_contains(&["web", "web_search"]); + let standalone_web_search = probe_with( |turn| { set_feature(turn, Feature::StandaloneWebSearch, /*enabled*/ true); diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 294b361f09..508021ff91 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -39,6 +39,14 @@ impl ToolExposure { pub fn is_direct(self) -> bool { matches!(self, Self::Direct | Self::DirectModelOnly) } + + /// Returns whether this tool can participate in code mode. + pub fn is_available_in_code_mode(self) -> bool { + match self { + Self::Direct | Self::Deferred => true, + Self::DirectModelOnly | Self::Hidden => false, + } + } } /// Shared runtime contract for model-visible tools.