From 0182ff34802c068815d5e0990ee94ec2d080b4fe Mon Sep 17 00:00:00 2001 From: Abhinav Date: Thu, 27 Aug 2026 16:20:33 +0000 Subject: [PATCH] Finalize model-specific tool plans in `ToolRouter` (#41195) ## Why Planning tools for a candidate model must not change the selected model or its Responses Lite tool inventory. Otherwise, preparing a fallback can overwrite metadata before the current model's request is sent. ## What changed - Build tool registries and visible specs from an explicit `ModelInfo`. - Store the finalized tool mode, code-mode names, namespace inventory, and terminal and child-management capabilities on `ToolRouter`. - Publish namespace inventory only after selecting the step's router, and use that router's mode when dispatching and recording tool calls. ## Testing Added coverage for candidate-model isolation, fallback compaction metadata, code-mode worker selection, normalized tool names, and router capabilities. GitOrigin-RevId: 56df2de8f42c7e93a4c4dcd117720e40a0cca3db --- codex-rs/core/src/session/mod.rs | 17 ++ codex-rs/core/src/session/tests.rs | 9 + codex-rs/core/src/session/turn.rs | 4 +- codex-rs/core/src/session/turn_context.rs | 4 +- .../core/src/stream_events_utils_tests.rs | 1 + codex-rs/core/src/tools/code_mode/mod.rs | 42 +++- codex-rs/core/src/tools/mod.rs | 9 +- codex-rs/core/src/tools/parallel.rs | 12 +- codex-rs/core/src/tools/router.rs | 104 ++++++++- codex-rs/core/src/tools/router_tests.rs | 3 + codex-rs/core/src/tools/spec_plan.rs | 217 ++++++++++++------ codex-rs/core/src/tools/spec_plan_tests.rs | 215 +++++++++++++---- codex-rs/core/tests/suite/compact.rs | 22 ++ 13 files changed, 532 insertions(+), 127 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 7a87eb97b6..9867606a20 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3398,6 +3398,8 @@ impl Session { let tool_router = turn::built_tools( self.as_ref(), turn_context.as_ref(), + // TODO(CDXENT-441): use the step scoped model + turn_context.model_info(), &environments, &mcp, &extension_data, @@ -3405,6 +3407,21 @@ impl Session { ) .or_cancel(cancellation_token) .await??; + // Publish inventory after planning rather than during finalization, so constructing + // additional candidate plans cannot overwrite turn-wide metadata. + if turn_context + .config + .tool_registry + .turn_metadata_includes_tool_info + && turn_context.model_info().use_responses_lite + { + turn_context.turn_metadata_state.set_tool_namespaces_info( + tool_router + .tool_namespaces_info() + .cloned() + .unwrap_or_default(), + ); + } Ok(Arc::new(StepContext { settings, token_budget, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ecfe7b39e9..a42d748a9c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -70,6 +70,7 @@ use codex_protocol::models::ImageDetail; use codex_protocol::models::PermissionProfile; use codex_protocol::models::SandboxEnforcement; use codex_protocol::openai_models::ModelServiceTier; +use codex_protocol::openai_models::ToolMode; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; @@ -257,6 +258,10 @@ impl StepContext { tool_router: Arc::new(ToolRouter::from_parts( ToolRegistry::empty_for_test(), Vec::new(), + ToolMode::Direct, + BTreeMap::new(), + /*tool_namespaces_info*/ None, + &[], )), loaded_agents_md: None, }) @@ -877,6 +882,7 @@ pub(crate) fn tool_registry_for_test_step( ) -> (ToolRegistry, Vec) { let mut registry = crate::tools::spec_plan::build_core_tool_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -884,6 +890,7 @@ pub(crate) fn tool_registry_for_test_step( ); let hosted_specs = crate::tools::spec_plan::append_source_tools( step_context.turn.as_ref(), + step_context.turn.model_info(), &mut registry, Vec::new(), Vec::new(), @@ -897,6 +904,7 @@ fn test_tool_runtime(session: Arc, turn_context: Arc) -> T 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(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), @@ -11612,6 +11620,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { let (registry, hosted_specs) = tool_registry_for_test_step(step_context.as_ref()); let router = ToolRouter::from_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index bdc542fea8..4a749b28b8 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1509,13 +1509,14 @@ pub(crate) async fn prepare_tool_recommendations( skip_all, fields( turn_id = %turn_context.sub_id, - model = %turn_context.model_info().slug, + model = %model_info.slug, apps_enabled = turn_context.apps_enabled() ) )] pub(crate) async fn built_tools( sess: &Session, turn_context: &TurnContext, + model_info: &codex_protocol::openai_models::ModelInfo, environments: &TurnEnvironmentSnapshot, mcp: &Arc, step_store: &ExtensionData, @@ -1584,6 +1585,7 @@ pub(crate) async fn built_tools( Ok(Arc::new(build_tool_router( sess, turn_context, + model_info, environments, mcp, apps_enabled, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 9ce1429f7d..93c9732f25 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1046,14 +1046,14 @@ impl Session { if !tc.code_mode_available && matches!( - crate::tools::requested_tool_mode(tc), + crate::tools::requested_tool_mode(tc, tc.model_info()), codex_protocol::openai_models::ToolMode::CodeMode | codex_protocol::openai_models::ToolMode::CodeModeOnly ) && let Some(message) = self .services .code_mode_service - .take_unavailable_warning(crate::tools::effective_tool_mode(tc)) + .take_unavailable_warning(crate::tools::effective_tool_mode(tc, tc.model_info())) { self.send_event(tc, EventMsg::Warning(WarningEvent { message })) .await; diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 20893b7fdd..3ccaea092e 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -291,6 +291,7 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { 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(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 6cbc1a2d1c..1465766a82 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -34,7 +34,6 @@ use crate::tools::ExecutedToolCallRecorder; use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; -use crate::tools::effective_tool_mode; use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; @@ -205,8 +204,7 @@ impl CodeModeService { tracker: SharedTurnDiffTracker, ) -> Option { let turn = &step_context.turn; - let tool_mode = effective_tool_mode(turn); - if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { + if !step_context.tool_router.requires_code_mode_worker() { return None; } @@ -442,14 +440,52 @@ fn build_freeform_tool_payload( #[cfg(test)] mod tests { + use std::collections::BTreeMap; + use std::sync::Arc; + use super::build_nested_tool_payload; use super::truncate_code_mode_result; + use crate::session::step_context::StepContext; + use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; + use crate::tools::registry::ToolRegistry; + use crate::tools::router::ToolRouter; + use crate::turn_diff_tracker::TurnDiffTracker; use codex_code_mode::CodeModeToolKind; use codex_protocol::models::FunctionCallOutputContentItem; + use codex_protocol::openai_models::ToolMode; use codex_tools::ToolName; use serde_json::json; + #[tokio::test] + async fn turn_worker_uses_step_router_mode_instead_of_admitted_turn() { + let (session, turn) = make_session_and_context().await; + assert_eq!( + crate::tools::effective_tool_mode(&turn, turn.model_info()), + ToolMode::Direct + ); + let session = Arc::new(session); + let step_context = StepContext::for_test(Arc::new(turn)); + let router = Arc::new(ToolRouter::from_parts( + ToolRegistry::empty_for_test(), + Vec::new(), + ToolMode::CodeModeOnly, + BTreeMap::new(), + /*tool_namespaces_info*/ None, + &[], + )); + let step_context = step_context.with_tool_router_for_test(router); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + + let worker = + session + .services + .code_mode_service + .start_turn_worker(&session, step_context, tracker); + + assert!(worker.is_some()); + } + #[test] fn build_nested_tool_payload_uses_function_kind() { let payload = build_nested_tool_payload( diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index e262e00cf0..ff1b6d063d 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -25,6 +25,7 @@ use crate::session::turn_context::TurnContext; pub(crate) use approvals::ApprovalContext; use codex_features::Feature; use codex_protocol::exec_output::ExecToolCallOutput; +use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ToolMode; use codex_tools::ToolName; use codex_utils_output_truncation::TruncationPolicy; @@ -64,8 +65,8 @@ pub(crate) fn tool_user_shell_type( } } -pub(crate) fn requested_tool_mode(turn_context: &TurnContext) -> ToolMode { - turn_context.model_info().tool_mode.unwrap_or_else(|| { +pub(crate) fn requested_tool_mode(turn_context: &TurnContext, model_info: &ModelInfo) -> ToolMode { + model_info.tool_mode.unwrap_or_else(|| { if turn_context.config.features.enabled(Feature::CodeModeOnly) { ToolMode::CodeModeOnly } else if turn_context.config.features.enabled(Feature::CodeMode) { @@ -76,8 +77,8 @@ pub(crate) fn requested_tool_mode(turn_context: &TurnContext) -> ToolMode { }) } -pub(crate) fn effective_tool_mode(turn_context: &TurnContext) -> ToolMode { - let requested_tool_mode = requested_tool_mode(turn_context); +pub(crate) fn effective_tool_mode(turn_context: &TurnContext, model_info: &ModelInfo) -> ToolMode { + let requested_tool_mode = requested_tool_mode(turn_context, model_info); if !turn_context.code_mode_available && requested_tool_mode == ToolMode::CodeMode && !turn_context.config.code_mode.disable_in_process_fallback diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 7886e4ee98..6b92c89270 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -106,7 +106,7 @@ impl ToolCallRuntime { executed_tool_calls.record_tool_call( &call, &source, - super::effective_tool_mode(&self.step_context.turn), + self.step_context.tool_router.tool_mode(), ); } let router = &self.step_context.tool_router; @@ -349,6 +349,7 @@ impl Drop for ToolCallTimingGuard { #[cfg(test)] mod tests { use super::*; + use std::collections::BTreeMap; use std::time::Duration; use crate::session::step_context::StepContext; @@ -362,6 +363,7 @@ mod tests { use codex_extension_api::ToolCallOutcome; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputPayload; + use codex_protocol::openai_models::ToolMode; use pretty_assertions::assert_eq; use tokio::sync::Notify; use tokio::sync::oneshot; @@ -425,6 +427,10 @@ mod tests { let router = Arc::new(ToolRouter::from_parts( ToolRegistry::from_tools([handler]), Vec::new(), + ToolMode::Direct, + BTreeMap::new(), + /*tool_namespaces_info*/ None, + &[], )); let step_context = step_context.with_tool_router_for_test(router); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); @@ -610,6 +616,10 @@ mod tests { let router = Arc::new(ToolRouter::from_parts( ToolRegistry::from_tools([handler]), Vec::new(), + ToolMode::Direct, + BTreeMap::new(), + /*tool_namespaces_info*/ None, + &[], )); let step_context = step_context.with_tool_router_for_test(router); 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 e5e6124b44..f7d65f5185 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,4 +1,5 @@ use crate::function_tool::FunctionCallError; +use crate::responses_metadata::TurnToolNamespacesInfo; use crate::session::session::Session; use crate::session::step_context::StepContext; #[cfg(test)] @@ -16,7 +17,11 @@ use crate::tools::registry::ToolRegistry; use crate::tools::spec_plan::finalize_tool_router; use codex_protocol::models::ResponseItem; use codex_protocol::models::SearchToolCallParams; +#[cfg(test)] +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ToolMode; use codex_tools::DiscoverableTool; +use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; use codex_tools::ToolSpec; use std::borrow::Cow; @@ -65,9 +70,14 @@ pub(crate) fn tool_log_payload<'a>( payload.log_payload() } +/// One finalized tool plan: its advertised surfaces and matching executable runtimes. pub struct ToolRouter { registry: ToolRegistry, model_visible_specs: Arc<[ToolSpec]>, + tool_mode: ToolMode, + code_mode_tool_names: BTreeMap, + tool_namespaces_info: Option, + can_manage_children: bool, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -86,12 +96,14 @@ impl ToolRouter { #[cfg(test)] pub(crate) fn from_registry( turn_context: &TurnContext, + model_info: &ModelInfo, registry: ToolRegistry, hosted_specs: Vec, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { finalize_tool_router( turn_context, + model_info, registry, hosted_specs, tool_search_handler_cache, @@ -99,17 +111,103 @@ impl ToolRouter { .expect("test tool registry should not contain duplicate tools") } - pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec) -> Self { - Self { + pub(crate) fn from_parts( + registry: ToolRegistry, + model_visible_specs: Vec, + tool_mode: ToolMode, + code_mode_tool_names: BTreeMap, + tool_namespaces_info: Option, + child_management_tools: &[ToolName], + ) -> Self { + let mut router = Self { registry, model_visible_specs: model_visible_specs.into(), - } + tool_mode, + code_mode_tool_names, + tool_namespaces_info, + can_manage_children: false, + }; + router.can_manage_children = !child_management_tools.is_empty() + && child_management_tools + .iter() + .all(|name| router.exposes_tool(name)); + router } pub(crate) fn model_visible_specs(&self) -> Arc<[ToolSpec]> { Arc::clone(&self.model_visible_specs) } + pub(crate) fn tool_mode(&self) -> ToolMode { + self.tool_mode + } + + /// Code Mode still needs its dispatcher when the nested tool set is empty. + pub(crate) fn requires_code_mode_worker(&self) -> bool { + matches!(self.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) + } + + /// The normalized nested identities chosen after exclusions and collisions. + // Consumed by the follow-up cell-origin migration. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn code_mode_tool_names(&self) -> &BTreeMap { + &self.code_mode_tool_names + } + + /// Optional request inventory for this exact plan, without publishing it to turn state. + pub(crate) fn tool_namespaces_info(&self) -> Option<&TurnToolNamespacesInfo> { + self.tool_namespaces_info.as_ref() + } + + /// Whether the model can both start and interact with a terminal process. + // Consumed by the follow-up live tool-plan selection. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn has_terminal_controls(&self) -> bool { + self.exposes_tool(&ToolName::plain("exec_command")) + && self.exposes_tool(&ToolName::plain("write_stdin")) + } + + /// Whether the configured collaboration backend's child-management tools remain exposed. + // Consumed by the follow-up live tool-plan selection. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn can_manage_children(&self) -> bool { + self.can_manage_children + } + + // Answers if the tool plan lets the model invoke the tool directly, through code mode, or deferred tool search. + fn exposes_tool(&self, name: &ToolName) -> bool { + let name = name.clone().with_default_namespace(); + if self + .code_mode_tool_names + .values() + .any(|nested| nested.clone().with_default_namespace() == name) + || self.model_visible_specs.iter().any(|spec| match spec { + ToolSpec::Function(_) | ToolSpec::Freeform(_) => { + name.is_default_namespace() && spec.name() == name.name + } + ToolSpec::Namespace(namespace) => { + name.namespace.as_deref() == Some(namespace.name.as_str()) + && namespace.tools.iter().any(|tool| match tool { + ResponsesApiNamespaceTool::Function(tool) => tool.name == name.name, + ResponsesApiNamespaceTool::Custom(tool) => tool.name == name.name, + }) + } + ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => false, + }) + { + return true; + } + self.model_visible_specs + .iter() + .any(|spec| matches!(spec, ToolSpec::ToolSearch { .. })) + && self.registry.entries().any(|tool| { + tool.exposure.is_deferred() + && tool.runtime.tool_name().with_default_namespace() == name + && (tool.runtime.immutable_spec().is_some() + || tool.runtime.search_info().is_some()) + }) + } + pub(crate) fn deferred_tool_namespaces(&self) -> BTreeMap { self.registry.deferred_tool_namespaces() } diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 7a12768805..d01de04d99 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -140,6 +140,7 @@ fn test_tool_router( ) -> ToolRouter { let mut registry = build_core_tool_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -147,6 +148,7 @@ fn test_tool_router( ); let hosted_specs = append_source_tools( step_context.turn.as_ref(), + step_context.turn.model_info(), &mut registry, mcp_tools, extension_tool_executors, @@ -154,6 +156,7 @@ fn test_tool_router( ); ToolRouter::from_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 9e50225a29..9d789113ad 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -40,6 +40,7 @@ use crate::tools::handlers::multi_agents::WaitAgentHandler; use crate::tools::handlers::multi_agents_common::DEFAULT_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents_common::MAX_WAIT_TIMEOUT_MS; use crate::tools::handlers::multi_agents_common::MIN_WAIT_TIMEOUT_MS; +use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; use crate::tools::handlers::multi_agents_spec::SpawnAgentToolOptions; use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions; use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2; @@ -72,6 +73,7 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; +use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::MultiAgentVersion; use codex_tools::ResponsesApiNamespace; @@ -105,6 +107,7 @@ const IMAGEGEN_TOOL_NAME: &str = "imagegen"; #[derive(Clone, Copy)] struct CoreToolPlanContext<'a> { turn_context: &'a TurnContext, + model_info: &'a ModelInfo, environments: &'a TurnEnvironmentSnapshot, mcp: &'a codex_mcp::McpBinding, tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, @@ -114,9 +117,11 @@ struct CoreToolPlanContext<'a> { } #[instrument(level = "trace", skip_all)] +#[allow(clippy::too_many_arguments)] pub(crate) fn build_tool_router( session: &Session, turn_context: &TurnContext, + model_info: &ModelInfo, environments: &TurnEnvironmentSnapshot, mcp: &Arc, apps_enabled: bool, @@ -131,6 +136,7 @@ pub(crate) fn build_tool_router( .get::(); let context = CoreToolPlanContext { turn_context, + model_info, environments, mcp, tool_suggest_candidates, @@ -149,6 +155,7 @@ pub(crate) fn build_tool_router( { append_extension_tool_executors( turn_context, + model_info, history_tools.0.iter().cloned(), &mut registry, ); @@ -160,21 +167,33 @@ pub(crate) fn build_tool_router( &turn_context.config, apps_enabled, &mcp.config().mcp_server_catalog, - search_tool_enabled(turn_context), + search_tool_enabled(turn_context, model_info), + &mut registry, + ); + apply_mcp_tool_exposure_policy( + turn_context, + model_info, + mcp, + ®istered_mcp_tools, &mut registry, ); - apply_mcp_tool_exposure_policy(turn_context, mcp, ®istered_mcp_tools, &mut registry); let standalone_web_search_tool = append_extension_tool_executors( turn_context, + model_info, extension_tool_executors(session, step_store), &mut registry, ); append_dynamic_tool_runtimes(&turn_context.dynamic_tools, &mut registry); - hosted_model_tool_specs(turn_context, standalone_web_search_tool.as_slice()) + hosted_model_tool_specs( + turn_context, + model_info, + standalone_web_search_tool.as_slice(), + ) }; finalize_tool_router( turn_context, + model_info, registry, hosted_specs, &session.services.tool_search_handler_cache, @@ -183,6 +202,7 @@ pub(crate) fn build_tool_router( fn apply_mcp_tool_exposure_policy( turn_context: &TurnContext, + model_info: &ModelInfo, mcp: &codex_mcp::McpBinding, registered_mcp_tools: &HashSet, registry: &mut ToolRegistry, @@ -228,9 +248,9 @@ fn apply_mcp_tool_exposure_policy( exposures = exposures.difference(ToolExposures::DEFERRED | ToolExposures::CODE_MODE); } - exposures = if search_tool_enabled(turn_context) + exposures = if search_tool_enabled(turn_context, model_info) && exposures.contains(ToolExposures::DEFERRED) - && (effective_tool_mode(turn_context) != ToolMode::CodeModeOnly + && (effective_tool_mode(turn_context, model_info) != ToolMode::CodeModeOnly || exposures.contains(ToolExposures::CODE_MODE)) { exposures.difference(ToolExposures::DIRECT) @@ -258,6 +278,7 @@ fn apply_mcp_tool_exposure_policy( #[instrument(level = "trace", skip_all)] pub(crate) fn build_core_tool_registry( turn_context: &TurnContext, + model_info: &ModelInfo, environments: &TurnEnvironmentSnapshot, mcp: &codex_mcp::McpBinding, tool_suggest_candidates: Option<&crate::tools::router::ToolSuggestCandidates>, @@ -267,6 +288,7 @@ pub(crate) fn build_core_tool_registry( crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { turn_context, + model_info, environments, mcp, tool_suggest_candidates, @@ -283,6 +305,7 @@ pub(crate) fn build_core_tool_registry( #[instrument(level = "trace", skip_all)] pub(crate) fn append_source_tools( turn_context: &TurnContext, + model_info: &ModelInfo, registry: &mut ToolRegistry, mcp_tools: Vec, extension_tool_executors: impl IntoIterator< @@ -297,10 +320,18 @@ pub(crate) fn append_source_tools( 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, registry); + let standalone_web_search_tool = append_extension_tool_executors( + turn_context, + model_info, + extension_tool_executors, + registry, + ); append_dynamic_tool_runtimes(dynamic_tools, registry); - hosted_model_tool_specs(turn_context, standalone_web_search_tool.as_slice()) + hosted_model_tool_specs( + turn_context, + model_info, + standalone_web_search_tool.as_slice(), + ) } #[instrument(level = "trace", skip_all)] @@ -325,15 +356,14 @@ pub(crate) fn extension_tool_executors<'a>( #[instrument(level = "trace", skip_all)] pub(crate) fn finalize_tool_router( turn_context: &TurnContext, + model_info: &ModelInfo, mut registry: ToolRegistry, hosted_specs: Vec, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> CodexResult { apply_direct_model_only_namespace_overrides(turn_context, &mut registry); - let code_mode_enabled = matches!( - effective_tool_mode(turn_context), - ToolMode::CodeMode | ToolMode::CodeModeOnly - ); + let tool_mode = effective_tool_mode(turn_context, model_info); + let code_mode_enabled = matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly); if code_mode_enabled { for tool_name in [ ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME), @@ -345,7 +375,7 @@ pub(crate) fn finalize_tool_router( } } let tool_search_name = ToolName::plain(TOOL_SEARCH_TOOL_NAME); - if search_tool_enabled(turn_context) + if search_tool_enabled(turn_context, model_info) && registry.entries().any(|tool| { tool.runtime.tool_name() != tool_search_name && tool.exposure.is_deferred() @@ -381,12 +411,13 @@ pub(crate) fn finalize_tool_router( append_tool_search_executor(turn_context, &mut registry, tool_search_handler_cache); } - let code_mode_tool_names = register_code_mode_executors(turn_context, &mut registry); + let code_mode_tool_names = + register_code_mode_executors(turn_context, model_info, &mut registry); let include_tool_namespaces_info = turn_context .config .tool_registry .turn_metadata_includes_tool_info - && turn_context.model_info().use_responses_lite; + && model_info.use_responses_lite; if turn_context.config.tool_registry.error_on_tool_collisions { if let Some(tool_name) = registry.first_collision() { @@ -444,19 +475,28 @@ pub(crate) fn finalize_tool_router( } } - let model_visible_specs = - build_model_visible_specs(turn_context, ®istry, &code_mode_tool_names, hosted_specs); - if include_tool_namespaces_info { - turn_context - .turn_metadata_state - .set_tool_namespaces_info(collect_tool_namespaces_info( - ®istry, - &code_mode_tool_names, - &model_visible_specs, - )); - } + let model_visible_specs = build_model_visible_specs( + turn_context, + model_info, + ®istry, + &code_mode_tool_names, + hosted_specs, + ); + let tool_namespaces_info = include_tool_namespaces_info + .then(|| { + collect_tool_namespaces_info(®istry, &code_mode_tool_names, &model_visible_specs) + }) + .filter(|info| !info.is_empty()); + let child_management_tools = required_child_management_tool_names(turn_context, model_info); - Ok(ToolRouter::from_parts(registry, model_visible_specs)) + Ok(ToolRouter::from_parts( + registry, + model_visible_specs, + tool_mode, + code_mode_tool_names, + tool_namespaces_info, + &child_management_tools, + )) } fn apply_direct_model_only_namespace_overrides( @@ -495,6 +535,7 @@ fn apply_direct_model_only_namespace_overrides( #[instrument(level = "trace", skip_all)] fn build_model_visible_specs( turn_context: &TurnContext, + model_info: &ModelInfo, registry: &ToolRegistry, code_mode_tool_names: &BTreeMap, hosted_specs: Vec, @@ -507,13 +548,14 @@ fn build_model_visible_specs( } let tool_name = tool.runtime.tool_name(); - if is_hidden_by_code_mode_only(turn_context, &tool_name, exposure) { + if is_hidden_by_code_mode_only(turn_context, model_info, &tool_name, exposure) { continue; } let spec = tool.runtime.spec(); specs.push(spec_for_model_request( turn_context, + model_info, exposure, &tool_name, code_mode_tool_names, @@ -532,12 +574,13 @@ fn build_model_visible_specs( fn spec_for_model_request( turn_context: &TurnContext, + model_info: &ModelInfo, exposure: ToolExposure, tool_name: &ToolName, code_mode_tool_names: &BTreeMap, spec: ToolSpec, ) -> ToolSpec { - let tool_mode = effective_tool_mode(turn_context); + let tool_mode = effective_tool_mode(turn_context, model_info); if matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) && exposure.is_available_in_code_mode() && !is_excluded_from_code_mode(turn_context, tool_name) @@ -557,17 +600,18 @@ fn spec_for_model_request( #[instrument(level = "trace", skip_all)] fn hosted_model_tool_specs( turn_context: &TurnContext, + model_info: &ModelInfo, registered_extension_tool_names: &[ToolName], ) -> Vec { // Responses Lite accepts schemas for client-executed tools, not hosted Responses tools. - if turn_context.model_info().use_responses_lite + if model_info.use_responses_lite || crate::guardian::is_basic_session_source(&turn_context.session_source) { return Vec::new(); } let mut specs = Vec::new(); - let standalone_web_search_available = standalone_web_search_enabled(turn_context) + let standalone_web_search_available = standalone_web_search_enabled(turn_context, model_info) && registered_extension_tool_names.contains(&ToolName::namespaced("web", "run")); // `Some(Cached/Live/Disabled)` are the options for mode when standalone search is unavailable // and the provider supports hosted search. `None` prevents emitting a hosted search tool. @@ -580,15 +624,15 @@ fn hosted_model_tool_specs( if let Some(hosted_web_search_tool) = create_web_search_tool(WebSearchToolOptions { web_search_mode, web_search_config, - web_search_tool_type: turn_context.model_info().web_search_tool_type, + web_search_tool_type: model_info.web_search_tool_type, }) { specs.push(hosted_web_search_tool); } specs } -pub(crate) fn search_tool_enabled(turn_context: &TurnContext) -> bool { - turn_context.model_info().supports_search_tool && namespace_tools_enabled(turn_context) +pub(crate) fn search_tool_enabled(turn_context: &TurnContext, model_info: &ModelInfo) -> bool { + model_info.supports_search_tool && namespace_tools_enabled(turn_context) } pub(crate) fn tool_suggest_enabled(turn_context: &TurnContext) -> bool { @@ -606,7 +650,7 @@ fn multi_agent_v2_enabled(turn_context: &TurnContext) -> bool { turn_context.multi_agent_version == MultiAgentVersion::V2 } -fn collab_tools_enabled(turn_context: &TurnContext) -> bool { +fn collab_tools_enabled(turn_context: &TurnContext, model_info: &ModelInfo) -> bool { match turn_context.multi_agent_version { MultiAgentVersion::Disabled => false, MultiAgentVersion::V1 => !exceeds_thread_spawn_depth_limit( @@ -615,12 +659,49 @@ fn collab_tools_enabled(turn_context: &TurnContext) -> bool { ), MultiAgentVersion::V2 => { turn_context.session_source.get_agent_path().is_none() - || turn_context.model_info().multi_agent_version == Some(MultiAgentVersion::V2) + || model_info.multi_agent_version == Some(MultiAgentVersion::V2) } } } -fn image_generation_available(turn_context: &TurnContext) -> bool { +fn required_child_management_tool_names( + turn_context: &TurnContext, + model_info: &ModelInfo, +) -> Vec { + if !collab_tools_enabled(turn_context, model_info) { + return Vec::new(); + } + + let (namespace, names): (_, &[&str]) = match turn_context.multi_agent_version { + MultiAgentVersion::Disabled => return Vec::new(), + MultiAgentVersion::V1 => ( + Some(MULTI_AGENT_V1_NAMESPACE), + &["send_input", "wait_agent", "resume_agent", "close_agent"], + ), + MultiAgentVersion::V2 => ( + namespace_tools_enabled(turn_context) + .then_some(turn_context.config.multi_agent_v2.tool_namespace.as_deref()) + .flatten(), + &[ + "send_message", + "followup_task", + "interrupt_agent", + "list_agents", + ], + ), + }; + let mut tools = names + .iter() + .map(|name| ToolName::new(namespace.map(str::to_owned), (*name).to_owned())) + .collect::>(); + if multi_agent_v2_enabled(turn_context) && turn_context.config.multi_agent_v2.wait_agent_enabled + { + tools.push(ToolName::new(namespace.map(str::to_owned), "wait_agent")); + } + tools +} + +fn image_generation_available(turn_context: &TurnContext, model_info: &ModelInfo) -> bool { if !turn_context .config .features @@ -645,11 +726,7 @@ fn image_generation_available(turn_context: &TurnContext) -> bool { return false; } - if !turn_context - .model_info() - .input_modalities - .contains(&InputModality::Image) - { + if !model_info.input_modalities.contains(&InputModality::Image) { return false; } @@ -693,10 +770,11 @@ fn agent_type_description( fn is_hidden_by_code_mode_only( turn_context: &TurnContext, + model_info: &ModelInfo, tool_name: &ToolName, exposure: ToolExposure, ) -> bool { - let tool_mode = effective_tool_mode(turn_context); + let tool_mode = effective_tool_mode(turn_context, model_info); tool_mode == ToolMode::CodeModeOnly && exposure.is_available_in_code_mode() && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name( @@ -717,9 +795,10 @@ fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) fn register_code_mode_executors( turn_context: &TurnContext, + model_info: &ModelInfo, registry: &mut ToolRegistry, ) -> BTreeMap { - let tool_mode = effective_tool_mode(turn_context); + let tool_mode = effective_tool_mode(turn_context, model_info); if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return BTreeMap::new(); } @@ -729,7 +808,7 @@ fn register_code_mode_executors( let mut exec_prompt_tool_specs = Vec::new(); let mut deferred_exec_prompt_tool_specs = Vec::new(); let mut included_deferred_mcp_output_schema = false; - let deferred_tools_guidance_enabled = search_tool_enabled(turn_context); + let deferred_tools_guidance_enabled = search_tool_enabled(turn_context, model_info); for tool in registry.entries() { let exposure = tool.exposure; if !exposure.is_available_in_code_mode() { @@ -807,10 +886,7 @@ fn register_code_mode_executors( &namespace_descriptions, turn_context.config.code_mode.default_exec_yield_time_ms, tool_mode == ToolMode::CodeModeOnly, - if unified_image_budget_enabled( - &turn_context.config.features, - turn_context.model_info(), - ) { + if unified_image_budget_enabled(&turn_context.config.features, model_info) { codex_code_mode::ImageDetailVisibility::Hidden } else { codex_code_mode::ImageDetailVisibility::Visible @@ -922,10 +998,7 @@ fn add_core_tool_sources(context: &CoreToolPlanContext<'_>, registry: &mut ToolR let include_environment_id = matches!(environment_mode, ToolEnvironmentMode::Multiple); if turn_context.config.features.enabled(Feature::ShellTool) && turn_context.config.features.enabled(Feature::UnifiedExec) - && !matches!( - turn_context.model_info().shell_type, - ConfigShellToolType::Disabled - ) + && !matches!(context.model_info.shell_type, ConfigShellToolType::Disabled) { registry.add(ExecCommandHandler::new(ExecCommandHandlerOptions { allow_login_shell: any_environment_allows_login_shell(context.environments), @@ -941,11 +1014,11 @@ fn add_core_tool_sources(context: &CoreToolPlanContext<'_>, registry: &mut ToolR if turn_context.config.features.enabled(Feature::ViewImage) { registry.add(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: can_request_original_image_detail( - turn_context.model_info(), + context.model_info, ), unified_image_budget: unified_image_budget_enabled( &turn_context.config.features, - turn_context.model_info(), + context.model_info, ), include_environment_id, })); @@ -960,10 +1033,10 @@ fn add_core_tool_sources(context: &CoreToolPlanContext<'_>, registry: &mut ToolR add_collaboration_tools(context, registry); } -fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool { +fn standalone_web_search_enabled(turn_context: &TurnContext, model_info: &ModelInfo) -> bool { namespace_tools_enabled(turn_context) && turn_context.provider.capabilities().web_search - && (turn_context.model_info().use_responses_lite + && (model_info.use_responses_lite || turn_context .config .features @@ -989,10 +1062,7 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, registry: &mut ToolRegistr if !environment_mode.has_environment() || !features.enabled(Feature::ShellTool) || !features.enabled(Feature::UnifiedExec) - || matches!( - turn_context.model_info().shell_type, - ConfigShellToolType::Disabled - ) + || matches!(context.model_info.shell_type, ConfigShellToolType::Disabled) { return; } @@ -1065,8 +1135,8 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool } if !turn_context.session_source.is_non_root_agent() - && turn_context - .model_info() + && context + .model_info .experimental_supported_tools .iter() .any(|tool| tool == "send_user_message_async") @@ -1111,15 +1181,13 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool )); } - if environment_mode.has_environment() - && turn_context.model_info().apply_patch_tool_type.is_some() - { + if environment_mode.has_environment() && context.model_info.apply_patch_tool_type.is_some() { let include_environment_id = matches!(environment_mode, ToolEnvironmentMode::Multiple); registry.add(ApplyPatchHandler::new(include_environment_id)); } - if turn_context - .model_info() + if context + .model_info .experimental_supported_tools .iter() .any(|tool| tool == "test_sync_tool") @@ -1131,11 +1199,11 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool let include_environment_id = matches!(environment_mode, ToolEnvironmentMode::Multiple); registry.add(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: can_request_original_image_detail( - turn_context.model_info(), + context.model_info, ), unified_image_budget: unified_image_budget_enabled( &turn_context.config.features, - turn_context.model_info(), + context.model_info, ), include_environment_id, })); @@ -1145,7 +1213,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool #[instrument(level = "trace", skip_all)] fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut ToolRegistry) { let turn_context = context.turn_context; - if collab_tools_enabled(turn_context) { + if collab_tools_enabled(turn_context, context.model_info) { if multi_agent_v2_enabled(turn_context) { let exposure = if turn_context.config.multi_agent_v2.non_code_mode_only { ToolExposure::DirectModelOnly @@ -1205,7 +1273,7 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, registry: &mut Too } else { let agent_type_description = agent_type_description(turn_context, context.default_agent_type_description); - let exposure = if search_tool_enabled(turn_context) { + let exposure = if search_tool_enabled(turn_context, context.model_info) { ToolExposure::Deferred } else { ToolExposure::Direct @@ -1285,10 +1353,11 @@ fn append_tool_search_executor( fn append_extension_tool_executors( turn_context: &TurnContext, + model_info: &ModelInfo, executors: impl IntoIterator ToolExecutor>>>, registry: &mut ToolRegistry, ) -> Option { - let standalone_web_search_enabled = standalone_web_search_enabled(turn_context); + let standalone_web_search_enabled = standalone_web_search_enabled(turn_context, model_info); let web_search_mode_on = turn_context.config.web_search_mode.value() != WebSearchMode::Disabled; let mut standalone_web_search_tool = None; @@ -1299,7 +1368,7 @@ fn append_extension_tool_executors( continue; } if tool_name == ToolName::namespaced(IMAGE_GEN_NAMESPACE, IMAGEGEN_TOOL_NAME) - && !image_generation_available(turn_context) + && !image_generation_available(turn_context, model_info) { continue; } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index a334cd4679..76cda520ad 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -20,6 +20,7 @@ use codex_protocol::error::CodexErrorDetails; use codex_protocol::openai_models::ApplyPatchToolType; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; +use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ToolMode; use codex_protocol::openai_models::WebSearchToolType; use codex_protocol::protocol::EnvironmentConfigState; @@ -46,6 +47,7 @@ use crate::config::CurrentTimeReminderConfig; use crate::environment_selection::TurnEnvironmentState; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::TurnToolFunctionInfo; +use crate::responses_metadata::TurnToolNamespacesInfo; use crate::responses_metadata::TurnToolSource; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; @@ -74,12 +76,19 @@ struct ToolPlanInputs { dynamic_tools: Vec, } +#[derive(Debug, PartialEq)] struct ToolPlanProbe { visible_specs: Vec, visible_names: Vec, namespace_functions: BTreeMap>, registered_names: Vec, exposures: BTreeMap, + tool_namespaces_info: Option, + code_mode_tool_names: BTreeMap, + tool_mode: ToolMode, + requires_code_mode_worker: bool, + has_terminal_controls: bool, + can_manage_children: bool, } impl ToolPlanProbe { @@ -129,6 +138,12 @@ impl ToolPlanProbe { namespace_functions, registered_names, exposures, + tool_namespaces_info: router.tool_namespaces_info().cloned(), + code_mode_tool_names: router.code_mode_tool_names().clone(), + tool_mode: router.tool_mode(), + requires_code_mode_worker: router.requires_code_mode_worker(), + has_terminal_controls: router.has_terminal_controls(), + can_manage_children: router.can_manage_children(), } } @@ -204,29 +219,38 @@ async fn probe_with( ) -> ToolPlanProbe { let (_session, mut turn) = make_session_and_context().await; configure_turn(&mut turn); - let turn = Arc::new(turn); - let step_context = StepContext::for_test(Arc::clone(&turn)); + ToolPlanProbe::from_router(plan_with_model(&turn, turn.model_info(), inputs)) +} + +fn plan_with_model( + turn: &TurnContext, + model_info: &ModelInfo, + inputs: ToolPlanInputs, +) -> ToolRouter { + let mcp = codex_mcp::McpBinding::empty(mcp_config_for_test(&turn.config)); let mut registry = build_core_tool_registry( - step_context.turn.as_ref(), - &step_context.environments, - step_context.mcp.as_ref(), + turn, + model_info, + &turn.environments, + &mcp, inputs.tool_suggest_candidates.as_ref(), inputs.wait_for_environment_tool_config.as_ref(), ); let hosted_specs = append_source_tools( - step_context.turn.as_ref(), + turn, + model_info, &mut registry, inputs.tool_runtimes, inputs.extension_tool_executors, &inputs.dynamic_tools, ); - let router = ToolRouter::from_registry( - step_context.turn.as_ref(), + ToolRouter::from_registry( + turn, + model_info, registry, hosted_specs, &Default::default(), - ); - ToolPlanProbe::from_router(router) + ) } async fn probe(configure_turn: impl FnOnce(&mut TurnContext)) -> ToolPlanProbe { @@ -497,6 +521,7 @@ async fn internal_guardian_sessions_exclude_optional_core_tools() { let router = super::build_tool_router( &session, step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, &step_context.mcp, /*apps_enabled*/ false, @@ -548,6 +573,7 @@ async fn internal_guardian_sessions_respect_managed_shell_restrictions() { let router = super::build_tool_router( &session, step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, &step_context.mcp, /*apps_enabled*/ false, @@ -583,6 +609,7 @@ async fn internal_guardian_sessions_preserve_code_mode() { let router = super::build_tool_router( &session, step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, &step_context.mcp, /*apps_enabled*/ false, @@ -652,6 +679,7 @@ async fn internal_guardian_sessions_require_managed_secondary_environments() { let router = super::build_tool_router( &session, step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, &step_context.mcp, /*apps_enabled*/ false, @@ -867,6 +895,7 @@ async fn shell_family_registers_only_unified_exec_tools() { plan.assert_visible_contains(&["exec_command", "write_stdin"]); plan.assert_registered_contains(&["exec_command", "write_stdin"]); + assert!(plan.has_terminal_controls); assert!(has_parameter(plan.visible_spec("exec_command"), "shell")); } @@ -907,6 +936,7 @@ async fn login_shell_parameter_follows_selected_environment() { has_parameter(plan.visible_spec("exec_command"), "login"), allow_login_shell ); + assert!(plan.has_terminal_controls); } } } @@ -959,6 +989,7 @@ async fn disabling_shell_tools_disables_command_tools_for_all_environments() { .await; remote_environment.assert_visible_lacks(&["exec_command", "write_stdin"]); remote_environment.assert_registered_lacks(&["exec_command", "write_stdin"]); + assert!(!remote_environment.has_terminal_controls); let multiple_local_environments = probe(|turn| { set_feature(turn, Feature::ShellTool, /*enabled*/ false); @@ -1159,6 +1190,7 @@ async fn environment_count_controls_environment_backed_tools() { "view_image", "request_permissions", ]); + assert!(!no_environment.has_terminal_controls); let multiple_environments = probe(|turn| { duplicate_primary_environment(turn); @@ -1176,6 +1208,7 @@ async fn environment_count_controls_environment_backed_tools() { "view_image", "request_permissions", ]); + assert!(multiple_environments.has_terminal_controls); assert!(has_parameter( multiple_environments.visible_spec("exec_command"), "environment_id" @@ -1206,14 +1239,16 @@ async fn environment_tools_follow_the_step_context() { let plan = ToolPlanProbe::from_router(ToolRouter::from_registry( turn.as_ref(), + turn.model_info(), build_core_tool_registry( turn.as_ref(), + turn.model_info(), &environments, mcp.as_ref(), /*tool_suggest_candidates*/ None, /*wait_for_environment_tool_config*/ None, ), - super::hosted_model_tool_specs(turn.as_ref(), &[]), + super::hosted_model_tool_specs(turn.as_ref(), turn.model_info(), &[]), &Default::default(), )); @@ -1424,9 +1459,8 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() { #[tokio::test] async fn tool_namespaces_info_is_opt_in_and_tracks_mcp_exposure() { - for enabled in [false, true] { - let mut metadata_state = None; - probe_with( + for (enabled, use_responses_lite) in [(false, true), (true, false), (true, true)] { + let plan = probe_with( |turn| { update_config(turn, |config| { config.tool_registry.turn_metadata_includes_tool_info = enabled; @@ -1434,9 +1468,8 @@ async fn tool_namespaces_info_is_opt_in_and_tracks_mcp_exposure() { set_feature(turn, Feature::CodeMode, /*enabled*/ true); update_turn_settings_for_test(turn, |settings| { Arc::make_mut(&mut settings.model_info).supports_search_tool = true; - Arc::make_mut(&mut settings.model_info).use_responses_lite = true; + Arc::make_mut(&mut settings.model_info).use_responses_lite = use_responses_lite; }); - metadata_state = Some(Arc::clone(&turn.turn_metadata_state)); }, ToolPlanInputs { tool_runtimes: vec![ @@ -1453,18 +1486,14 @@ async fn tool_namespaces_info_is_opt_in_and_tracks_mcp_exposure() { ) .await; - let metadata = metadata_state - .expect("tool planning should capture the turn metadata") - .to_responses_metadata( - "installation".to_string(), - "window".to_string(), - CodexResponsesRequestKind::Turn, + let Some(namespaces) = plan.tool_namespaces_info else { + assert!( + !enabled || !use_responses_lite, + "opted-in Responses Lite planning should return namespaces" ); - let Some(namespaces) = metadata.tool_namespaces_info else { - assert!(!enabled, "opted-in tool planning should publish namespaces"); continue; }; - assert!(enabled, "tool namespaces require opt-in"); + assert!(enabled && use_responses_lite); let namespace = namespaces .get("mcp__registry") @@ -1497,6 +1526,80 @@ async fn tool_namespaces_info_is_opt_in_and_tracks_mcp_exposure() { } } +#[tokio::test] +async fn candidate_model_plan_leaves_selected_model_and_inventory_unchanged() { + let (_session, mut turn) = make_session_and_context().await; + set_features(&mut turn, &[Feature::ShellTool, Feature::UnifiedExec]); + update_config(&mut turn, |config| { + config.tool_registry.turn_metadata_includes_tool_info = true; + }); + update_turn_settings_for_test(&mut turn, |settings| { + let model = Arc::make_mut(&mut settings.model_info); + model.tool_mode = Some(ToolMode::Direct); + model.use_responses_lite = true; + model.shell_type = ConfigShellToolType::Disabled; + model.apply_patch_tool_type = None; + }); + let selected_model = Arc::clone(turn.model_info()); + let selected = ToolPlanProbe::from_router(plan_with_model( + &turn, + turn.model_info(), + ToolPlanInputs::default(), + )); + let selected_inventory = selected + .tool_namespaces_info + .clone() + .expect("selected plan inventory"); + turn.turn_metadata_state + .set_tool_namespaces_info(selected_inventory.clone()); + + let mut candidate_model = selected_model.as_ref().clone(); + candidate_model.tool_mode = Some(ToolMode::CodeModeOnly); + candidate_model.shell_type = ConfigShellToolType::UnifiedExec; + candidate_model.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + let candidate = ToolPlanProbe::from_router(plan_with_model( + &turn, + &candidate_model, + ToolPlanInputs::default(), + )); + + candidate.assert_visible_contains(&["exec", "wait"]); + candidate.assert_visible_lacks(&["exec_command", "write_stdin", "apply_patch"]); + candidate.assert_registered_contains(&["exec_command", "write_stdin", "apply_patch"]); + assert_eq!(candidate.tool_mode, ToolMode::CodeModeOnly); + assert!(candidate.requires_code_mode_worker); + assert!(candidate.has_terminal_controls); + assert_eq!( + candidate + .tool_namespaces_info + .as_ref() + .expect("candidate inventory")["functions"] + .functions["apply_patch"], + TurnToolFunctionInfo { + name: "apply_patch".to_string(), + direct: false, + code_mode_name: Some("apply_patch".to_string()), + deferred: false, + source: TurnToolSource::Harness, + } + ); + let metadata = turn.turn_metadata_state.to_responses_metadata( + "installation".to_string(), + "window".to_string(), + CodexResponsesRequestKind::Turn, + ); + assert_eq!(metadata.tool_namespaces_info, Some(selected_inventory)); + assert_eq!(turn.model_info(), &selected_model); + assert_eq!( + ToolPlanProbe::from_router(plan_with_model( + &turn, + turn.model_info(), + ToolPlanInputs::default(), + )), + selected, + ); +} + #[tokio::test] async fn strict_namespace_ownership_requires_tool_namespace_inventory_opt_in() { for (enabled, second_exposure) in [ @@ -1515,6 +1618,7 @@ async fn strict_namespace_ownership_requires_tool_namespace_inventory_opt_in() { let step_context = StepContext::for_test(Arc::new(turn)); let mut registry = build_core_tool_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -1536,6 +1640,7 @@ async fn strict_namespace_ownership_requires_tool_namespace_inventory_opt_in() { .collect(); let hosted_specs = append_source_tools( step_context.turn.as_ref(), + step_context.turn.model_info(), &mut registry, runtimes, Vec::new(), @@ -1543,6 +1648,7 @@ async fn strict_namespace_ownership_requires_tool_namespace_inventory_opt_in() { ); let result = super::finalize_tool_router( step_context.turn.as_ref(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), @@ -1772,6 +1878,7 @@ async fn strict_tool_collisions_reject_external_and_synthetic_duplicates() { let step_context = StepContext::for_test(Arc::clone(&turn)); let mut registry = build_core_tool_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, step_context.mcp.as_ref(), inputs.tool_suggest_candidates.as_ref(), @@ -1779,6 +1886,7 @@ async fn strict_tool_collisions_reject_external_and_synthetic_duplicates() { ); let hosted_specs = append_source_tools( step_context.turn.as_ref(), + step_context.turn.model_info(), &mut registry, inputs.tool_runtimes, inputs.extension_tool_executors, @@ -1787,6 +1895,7 @@ async fn strict_tool_collisions_reject_external_and_synthetic_duplicates() { let error = super::finalize_tool_router( step_context.turn.as_ref(), + step_context.turn.model_info(), registry, hosted_specs, &Default::default(), @@ -1915,7 +2024,6 @@ async fn code_mode_uses_the_first_normalized_tool_identity() { (true, ToolExposure::Direct, true), (true, ToolExposure::Deferred, false), ] { - let mut metadata_state = None; let plan = probe_with( |turn| { update_config(turn, |config| { @@ -1929,7 +2037,6 @@ async fn code_mode_uses_the_first_normalized_tool_identity() { Arc::make_mut(&mut settings.model_info).supports_search_tool = true; Arc::make_mut(&mut settings.model_info).use_responses_lite = true; }); - metadata_state = Some(Arc::clone(&turn.turn_metadata_state)); }, ToolPlanInputs { tool_runtimes: vec![mcp_runtime( @@ -1961,17 +2068,14 @@ async fn code_mode_uses_the_first_normalized_tool_identity() { }, ); - let metadata = metadata_state - .expect("tool planning should capture the turn metadata") - .to_responses_metadata( - "installation".to_string(), - "window".to_string(), - CodexResponsesRequestKind::Turn, - ); - let tool_namespaces = metadata + let tool_namespaces = plan .tool_namespaces_info .as_ref() - .expect("opted-in Responses Lite should receive the supported tools"); + .expect("opted-in Responses Lite plan should contain the supported tools"); + assert_eq!( + plan.code_mode_tool_names.get("normalized_alias__lookup"), + Some(&winner_name), + ); assert_eq!( tool_namespaces["normalized-alias"].functions["lookup"] .code_mode_name @@ -2040,6 +2144,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let first_step_context = StepContext::for_test(Arc::clone(&first_turn)); let mut first_registry = build_core_tool_registry( first_step_context.turn.as_ref(), + first_step_context.turn.model_info(), &first_step_context.environments, first_step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -2049,8 +2154,13 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { 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_step_context.turn.model_info(), first_registry, - super::hosted_model_tool_specs(first_step_context.turn.as_ref(), &[]), + super::hosted_model_tool_specs( + first_step_context.turn.as_ref(), + first_step_context.turn.model_info(), + &[], + ), &cache, ); let first_plan = ToolPlanProbe::from_router(first_router); @@ -2063,6 +2173,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let second_step_context = StepContext::for_test(Arc::clone(&second_turn)); let mut second_registry = build_core_tool_registry( second_step_context.turn.as_ref(), + second_step_context.turn.model_info(), &second_step_context.environments, second_step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -2072,8 +2183,13 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { 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_step_context.turn.model_info(), second_registry, - super::hosted_model_tool_specs(second_step_context.turn.as_ref(), &[]), + super::hosted_model_tool_specs( + second_step_context.turn.as_ref(), + second_step_context.turn.model_info(), + &[], + ), &cache, ); let second_plan = ToolPlanProbe::from_router(second_router); @@ -2117,6 +2233,7 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() { let step_context = StepContext::for_test(Arc::clone(&turn)); let mut registry = build_core_tool_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), &step_context.environments, step_context.mcp.as_ref(), /*tool_suggest_candidates*/ None, @@ -2131,8 +2248,13 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() { registry.register_external_with_exposure(tool.runtime, tool.exposure); let router = ToolRouter::from_registry( step_context.turn.as_ref(), + step_context.turn.model_info(), registry, - super::hosted_model_tool_specs(step_context.turn.as_ref(), &[]), + super::hosted_model_tool_specs( + step_context.turn.as_ref(), + step_context.turn.model_info(), + &[], + ), &cache, ); let plan = ToolPlanProbe::from_router(router); @@ -2303,6 +2425,10 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { codex_code_mode::PUBLIC_TOOL_NAME, codex_code_mode::WAIT_TOOL_NAME, ]); + assert_eq!( + (plain.tool_mode, plain.requires_code_mode_worker), + (ToolMode::Direct, false) + ); let code_mode_only = probe_with( |turn| { @@ -2322,6 +2448,13 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { codex_code_mode::PUBLIC_TOOL_NAME, codex_code_mode::WAIT_TOOL_NAME, ]); + assert_eq!( + ( + code_mode_only.tool_mode, + code_mode_only.requires_code_mode_worker + ), + (ToolMode::CodeModeOnly, true), + ); assert_eq!( code_mode_only.namespace_function_names("codex_app"), Vec::::new().as_slice() @@ -2484,6 +2617,7 @@ async fn multi_agent_feature_selects_one_agent_tool_family() { }) .await; v1.assert_visible_contains(&[MULTI_AGENT_V1_NAMESPACE]); + assert!(v1.can_manage_children); v1.assert_visible_lacks(&[ "spawn_agent", "send_input", @@ -2540,6 +2674,7 @@ async fn multi_agent_feature_selects_one_agent_tool_family() { }) .await; v2.assert_visible_contains(&[MULTI_AGENT_V2_NAMESPACE]); + assert!(v2.can_manage_children); v2.assert_visible_lacks(&[ "spawn_agent", "send_message", @@ -2674,6 +2809,7 @@ async fn multi_agent_v2_can_disable_wait_agent() { ); plan.assert_visible_lacks(&["clock"]); plan.assert_registered_lacks(&["collaboration.wait_agent", "clock.sleep"]); + assert!(plan.can_manage_children); } #[tokio::test] @@ -2856,6 +2992,7 @@ async fn multi_agent_v2_bedrock_workers_only_delegate_when_model_supports_v2() { let spawn_agent_name = ToolName::namespaced("agents", "spawn_agent").to_string(); let followup_task_name = ToolName::namespaced("agents", "followup_task").to_string(); + assert_eq!(plan.can_manage_children, supports_delegation); if supports_delegation { plan.assert_visible_contains(&["agents"]); plan.assert_registered_contains(&[&spawn_agent_name, &followup_task_name]); diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 40c74d0a04..bd97ff635e 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -2666,9 +2666,11 @@ async fn pre_sampling_compact_falls_back_after_previous_model_invalid_request_on let mut previous_model_info = model_info_with_context_window("gpt-5.4", /*context_window*/ 273_000); previous_model_info.slug = previous_model_family.to_string(); + previous_model_info.use_responses_lite = true; let mut next_model_info = model_info_with_context_window("gpt-5.4", /*context_window*/ 125_000); next_model_info.slug = next_model.to_string(); + next_model_info.use_responses_lite = false; let models_mock = mount_models_once( &server, @@ -2710,6 +2712,7 @@ async fn pre_sampling_compact_falls_back_after_previous_model_invalid_request_on .with_model(retired_model) .with_config(move |config| { config.model_provider = model_provider; + config.tool_registry.turn_metadata_includes_tool_info = true; set_test_compact_prompt(config); let _ = config.features.enable(Feature::RemoteCompactionV2); }); @@ -2755,6 +2758,25 @@ async fn pre_sampling_compact_falls_back_after_previous_model_invalid_request_on ); assert_eq!(requests[2].body_json()["model"].as_str(), Some(next_model)); assert_eq!(requests[3].body_json()["model"].as_str(), Some(next_model)); + + // Preparing the non-Lite fallback must not clear the previous model's inventory + // before its compaction attempt is sent. + let [first_metadata, compact_metadata] = [&requests[0], &requests[1]].map(|request| { + serde_json::from_str::( + request.body_json()["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("request should include turn metadata"), + ) + .expect("turn metadata should be valid JSON") + }); + let inventory = first_metadata["tool_namespaces_info"] + .as_object() + .expect("Responses Lite request should include tool inventory"); + assert!(!inventory.is_empty()); + assert_eq!( + compact_metadata["tool_namespaces_info"], + first_metadata["tool_namespaces_info"], + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)]