From 581291bf68ea593d40d19ab2f63754522edfcc0b Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 29 May 2026 00:43:23 -0700 Subject: [PATCH] [codex] Inline multi-agent version resolution --- codex-rs/core/src/lib.rs | 2 - codex-rs/core/src/multi_agent_version.rs | 63 ---------------------- codex-rs/core/src/session/review.rs | 24 +++++++-- codex-rs/core/src/session/turn_context.rs | 51 ++++++++++++------ codex-rs/core/src/tools/spec_plan.rs | 45 +++++++--------- codex-rs/core/src/tools/spec_plan_tests.rs | 32 +++++++---- 6 files changed, 96 insertions(+), 121 deletions(-) delete mode 100644 codex-rs/core/src/multi_agent_version.rs diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 239d38fbc7..aa5784d4fe 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -60,8 +60,6 @@ pub use codex_mcp::SandboxState; mod mcp_openai_file; mod mcp_tool_call; pub(crate) mod mention_syntax; -mod multi_agent_version; -mod tool_mode; pub(crate) mod utils; pub use mention_syntax::PLUGIN_TEXT_MENTION_SIGIL; pub use mention_syntax::TOOL_MENTION_SIGIL; diff --git a/codex-rs/core/src/multi_agent_version.rs b/codex-rs/core/src/multi_agent_version.rs deleted file mode 100644 index 245bce6fe2..0000000000 --- a/codex-rs/core/src/multi_agent_version.rs +++ /dev/null @@ -1,63 +0,0 @@ -use crate::config::ManagedFeatures; -use codex_features::Feature; -use codex_protocol::openai_models::ModelInfo; -use codex_protocol::openai_models::MultiAgentVersion; - -pub(crate) fn resolve_multi_agent_version( - model_info: &ModelInfo, - features: &ManagedFeatures, -) -> Option { - model_info - .multi_agent_version - .or_else(|| multi_agent_version_from_features(features)) -} - -fn multi_agent_version_from_features(features: &ManagedFeatures) -> Option { - if features.enabled(Feature::MultiAgentV2) { - Some(MultiAgentVersion::V2) - } else if features.enabled(Feature::Collab) { - Some(MultiAgentVersion::V1) - } else { - None - } -} - -#[cfg(test)] -mod tests { - use super::*; - use codex_features::Features; - - fn features(enabled: &[Feature]) -> ManagedFeatures { - let mut features = Features::default(); - for feature in enabled { - features.enable(*feature); - } - features.into() - } - - fn model_info(multi_agent_version: Option) -> ModelInfo { - let mut model_info = codex_models_manager::model_info::model_info_from_slug("test-model"); - model_info.multi_agent_version = multi_agent_version; - model_info - } - - #[test] - fn omitted_selector_follows_feature_flags() { - let multi_agent_version = resolve_multi_agent_version( - &model_info(/*multi_agent_version*/ None), - &features(&[Feature::MultiAgentV2]), - ); - - assert_eq!(multi_agent_version, Some(MultiAgentVersion::V2)); - } - - #[test] - fn explicit_selector_overrides_feature_flags() { - let multi_agent_version = resolve_multi_agent_version( - &model_info(Some(MultiAgentVersion::V1)), - &features(&[Feature::MultiAgentV2]), - ); - - assert_eq!(multi_agent_version, Some(MultiAgentVersion::V1)); - } -} diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index dd2382f445..aac7f2a9d0 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -1,6 +1,6 @@ use super::*; -use crate::multi_agent_version::resolve_multi_agent_version; -use crate::tool_mode::resolve_tool_mode; +use codex_protocol::openai_models::MultiAgentVersion; +use codex_protocol::openai_models::ToolMode; use std::sync::atomic::AtomicBool; /// Spawn a review thread using the given prompt. @@ -49,8 +49,24 @@ pub(super) async fn spawn_review_thread( let mut per_turn_config = (*config).clone(); per_turn_config.model = Some(model.clone()); per_turn_config.features = review_features.clone(); - let tool_mode = resolve_tool_mode(&model_info, &per_turn_config.features); - let multi_agent_version = resolve_multi_agent_version(&model_info, &per_turn_config.features); + let tool_mode = model_info.tool_mode.unwrap_or_else(|| { + if per_turn_config.features.enabled(Feature::CodeModeOnly) { + ToolMode::CodeModeOnly + } else if per_turn_config.features.enabled(Feature::CodeMode) { + ToolMode::CodeMode + } else { + ToolMode::Direct + } + }); + let multi_agent_version = model_info.multi_agent_version.or_else(|| { + if per_turn_config.features.enabled(Feature::MultiAgentV2) { + Some(MultiAgentVersion::V2) + } else if per_turn_config.features.enabled(Feature::Collab) { + Some(MultiAgentVersion::V1) + } else { + None + } + }); if let Err(err) = per_turn_config.web_search_mode.set(review_web_search_mode) { let fallback_value = per_turn_config.web_search_mode.value(); tracing::warn!( diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index ce3fd5ecaf..dd5f0618f7 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -2,8 +2,6 @@ use super::*; use crate::SkillLoadOutcome; use crate::config::GhostSnapshotConfig; use crate::environment_selection::ResolvedTurnEnvironments; -use crate::multi_agent_version::resolve_multi_agent_version; -use crate::tool_mode::resolve_tool_mode; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_protocol::SessionId; @@ -168,14 +166,6 @@ impl TurnContext { self.goal_tools_supported && self.features.get().enabled(Feature::Goals) } - pub(crate) fn code_mode_enabled(&self) -> bool { - matches!(self.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - } - - pub(crate) fn code_mode_only_enabled(&self) -> bool { - self.tool_mode == ToolMode::CodeModeOnly - } - pub(crate) async fn with_model( &self, model: String, @@ -186,8 +176,24 @@ impl TurnContext { let model_info = models_manager .get_model_info(model.as_str(), &config.to_models_manager_config()) .await; - let tool_mode = resolve_tool_mode(&model_info, &config.features); - let multi_agent_version = resolve_multi_agent_version(&model_info, &config.features); + let tool_mode = model_info.tool_mode.unwrap_or_else(|| { + if config.features.enabled(Feature::CodeModeOnly) { + ToolMode::CodeModeOnly + } else if config.features.enabled(Feature::CodeMode) { + ToolMode::CodeMode + } else { + ToolMode::Direct + } + }); + let multi_agent_version = model_info.multi_agent_version.or_else(|| { + if config.features.enabled(Feature::MultiAgentV2) { + Some(MultiAgentVersion::V2) + } else if config.features.enabled(Feature::Collab) { + Some(MultiAgentVersion::V1) + } else { + None + } + }); let truncation_policy = model_info.truncation_policy.into(); let supported_reasoning_levels = model_info .supported_reasoning_levels @@ -493,9 +499,24 @@ impl Session { ); let mut per_turn_config = per_turn_config; - let tool_mode = resolve_tool_mode(&model_info, &per_turn_config.features); - let multi_agent_version = - resolve_multi_agent_version(&model_info, &per_turn_config.features); + let tool_mode = model_info.tool_mode.unwrap_or_else(|| { + if per_turn_config.features.enabled(Feature::CodeModeOnly) { + ToolMode::CodeModeOnly + } else if per_turn_config.features.enabled(Feature::CodeMode) { + ToolMode::CodeMode + } else { + ToolMode::Direct + } + }); + let multi_agent_version = model_info.multi_agent_version.or_else(|| { + if per_turn_config.features.enabled(Feature::MultiAgentV2) { + Some(MultiAgentVersion::V2) + } else if per_turn_config.features.enabled(Feature::Collab) { + Some(MultiAgentVersion::V1) + } else { + None + } + }); per_turn_config.service_tier = get_service_tier( per_turn_config.service_tier, per_turn_config.features.enabled(Feature::FastMode), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index d22aa55d10..35f937dd80 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -61,6 +61,7 @@ use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::MultiAgentVersion; +use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_tools::DiscoverableTool; @@ -231,8 +232,10 @@ fn spec_for_model_request( exposure: ToolExposure, spec: ToolSpec, ) -> ToolSpec { - if code_mode_enabled(turn_context) - && exposure != ToolExposure::DirectModelOnly + if matches!( + turn_context.tool_mode, + ToolMode::CodeMode | ToolMode::CodeModeOnly + ) && exposure != ToolExposure::DirectModelOnly && codex_code_mode::is_code_mode_nested_tool(spec.name()) { codex_tools::augment_tool_spec_for_code_mode(spec) @@ -283,22 +286,6 @@ fn namespace_tools_enabled(turn_context: &TurnContext) -> bool { turn_context.provider.capabilities().namespace_tools } -fn code_mode_enabled(turn_context: &TurnContext) -> bool { - turn_context.code_mode_enabled() -} - -fn code_mode_only_enabled(turn_context: &TurnContext) -> bool { - turn_context.code_mode_only_enabled() -} - -fn multi_agent_v2_enabled(turn_context: &TurnContext) -> bool { - turn_context.multi_agent_version == Some(MultiAgentVersion::V2) -} - -fn collab_tools_enabled(turn_context: &TurnContext) -> bool { - turn_context.multi_agent_version.is_some() -} - fn goal_tools_enabled(turn_context: &TurnContext) -> bool { turn_context.goal_tools_enabled() && !matches!( @@ -357,7 +344,7 @@ fn standalone_image_generation_available( } fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOptions { - if multi_agent_v2_enabled(turn_context) { + if turn_context.multi_agent_version == Some(MultiAgentVersion::V2) { return WaitAgentTimeoutOptions { default_timeout_ms: turn_context.config.multi_agent_v2.default_wait_timeout_ms, min_timeout_ms: turn_context.config.multi_agent_v2.min_wait_timeout_ms, @@ -373,7 +360,7 @@ fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOpt } fn max_concurrent_threads_per_session(turn_context: &TurnContext) -> Option { - multi_agent_v2_enabled(turn_context).then_some( + (turn_context.multi_agent_version == Some(MultiAgentVersion::V2)).then_some( turn_context .config .multi_agent_v2 @@ -399,7 +386,7 @@ fn is_hidden_by_code_mode_only( tool_name: &ToolName, exposure: ToolExposure, ) -> bool { - code_mode_only_enabled(turn_context) + turn_context.tool_mode == ToolMode::CodeModeOnly && exposure != ToolExposure::DirectModelOnly && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name( tool_name, @@ -411,7 +398,10 @@ fn build_code_mode_executors( executors: &[Arc], deferred_tools_available: bool, ) -> Vec> { - if !code_mode_enabled(turn_context) { + if !matches!( + turn_context.tool_mode, + ToolMode::CodeMode | ToolMode::CodeModeOnly + ) { return vec![]; } @@ -445,7 +435,7 @@ fn build_code_mode_executors( create_code_mode_tool( &enabled_tools, &namespace_descriptions, - code_mode_only_enabled(turn_context), + turn_context.tool_mode == ToolMode::CodeModeOnly, deferred_tools_available, ), code_mode_nested_tool_specs, @@ -649,8 +639,8 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { let turn_context = context.turn_context; - if collab_tools_enabled(turn_context) { - if multi_agent_v2_enabled(turn_context) { + if turn_context.multi_agent_version.is_some() { + if turn_context.multi_agent_version == Some(MultiAgentVersion::V2) { let exposure = if turn_context.config.multi_agent_v2.non_code_mode_only { ToolExposure::DirectModelOnly } else { @@ -848,7 +838,10 @@ fn append_extension_tool_executors( .iter() .map(|executor| executor.tool_name()) .collect::>(); - if code_mode_enabled(turn_context) { + if matches!( + turn_context.tool_mode, + ToolMode::CodeMode | ToolMode::CodeModeOnly + ) { reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)); reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME)); } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 58802f6eb1..7a7061a774 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -31,10 +31,8 @@ use codex_tools::ToolSpec; use pretty_assertions::assert_eq; use serde_json::json; -use crate::multi_agent_version::resolve_multi_agent_version; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; -use crate::tool_mode::resolve_tool_mode; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; use crate::tools::router::ToolRouter; use crate::tools::router::ToolRouterParams; @@ -219,7 +217,24 @@ fn set_feature(turn: &mut TurnContext, feature: Feature, enabled: bool) { .expect("test feature should be disableable in config"); } turn.config = Arc::new(config); - resolve_model_selectors_for_turn(turn); + turn.tool_mode = turn.model_info.tool_mode.unwrap_or_else(|| { + if turn.config.features.enabled(Feature::CodeModeOnly) { + ToolMode::CodeModeOnly + } else if turn.config.features.enabled(Feature::CodeMode) { + ToolMode::CodeMode + } else { + ToolMode::Direct + } + }); + turn.multi_agent_version = turn.model_info.multi_agent_version.or_else(|| { + if turn.config.features.enabled(Feature::MultiAgentV2) { + Some(MultiAgentVersion::V2) + } else if turn.config.features.enabled(Feature::Collab) { + Some(MultiAgentVersion::V1) + } else { + None + } + }); } fn set_features(turn: &mut TurnContext, features: &[Feature]) { @@ -234,11 +249,6 @@ fn update_config(turn: &mut TurnContext, update: impl FnOnce(&mut crate::config: turn.config = Arc::new(config); } -fn resolve_model_selectors_for_turn(turn: &mut TurnContext) { - turn.tool_mode = resolve_tool_mode(&turn.model_info, &turn.config.features); - turn.multi_agent_version = resolve_multi_agent_version(&turn.model_info, &turn.config.features); -} - fn set_web_search_mode(turn: &mut TurnContext, mode: WebSearchMode) { update_config(turn, |config| { config @@ -812,7 +822,7 @@ async fn tool_mode_selector_overrides_feature_flags() { let direct = probe(|turn| { set_features(turn, &[Feature::CodeMode, Feature::CodeModeOnly]); turn.model_info.tool_mode = Some(ToolMode::Direct); - resolve_model_selectors_for_turn(turn); + turn.tool_mode = ToolMode::Direct; }) .await; direct.assert_visible_lacks(&[ @@ -826,7 +836,7 @@ async fn multi_agent_version_selector_overrides_feature_flags() { let v1 = probe(|turn| { set_feature(turn, Feature::MultiAgentV2, /*enabled*/ true); turn.model_info.multi_agent_version = Some(MultiAgentVersion::V1); - resolve_model_selectors_for_turn(turn); + turn.multi_agent_version = Some(MultiAgentVersion::V1); }) .await; @@ -843,7 +853,7 @@ async fn multi_agent_version_selector_overrides_feature_flags() { set_feature(turn, Feature::Collab, /*enabled*/ false); set_feature(turn, Feature::MultiAgentV2, /*enabled*/ false); turn.model_info.multi_agent_version = Some(MultiAgentVersion::V2); - resolve_model_selectors_for_turn(turn); + turn.multi_agent_version = Some(MultiAgentVersion::V2); }) .await;