[codex] Resolve multi-agent version on turn context

codex-rs/core/src/lib.rs

@@ -60,6 +60,7 @@ 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;

codex-rs/core/src/multi_agent_version.rs

@@ -0,0 +1,63 @@
+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<MultiAgentVersion> {
+    model_info
+        .multi_agent_version
+        .or_else(|| multi_agent_version_from_features(features))
+}
+
+fn multi_agent_version_from_features(features: &ManagedFeatures) -> Option<MultiAgentVersion> {
+    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<MultiAgentVersion>) -> 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));
+    }
+}

codex-rs/core/src/session/review.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::multi_agent_version::resolve_multi_agent_version;
 use crate::tool_mode::resolve_tool_mode;
 use std::sync::atomic::AtomicBool;
 
@@ -49,6 +50,7 @@ pub(super) async fn spawn_review_thread(
     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);
     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!(
@@ -99,6 +101,7 @@ pub(super) async fn spawn_review_thread(
         auth_manager: auth_manager_for_context,
         model_info: model_info.clone(),
         tool_mode,
+        multi_agent_version,
         session_telemetry: session_telemetry_for_context,
         provider: provider_for_context,
         reasoning_effort,

codex-rs/core/src/session/turn_context.rs

@@ -2,11 +2,13 @@ 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;
 use codex_protocol::models::AdditionalPermissionProfile;
+use codex_protocol::openai_models::MultiAgentVersion;
 use codex_protocol::openai_models::ToolMode;
 use codex_protocol::protocol::ThreadSource;
 use codex_protocol::protocol::TurnEnvironmentSelection;
@@ -58,6 +60,7 @@ pub struct TurnContext {
     pub(crate) auth_manager: Option<Arc<AuthManager>>,
     pub(crate) model_info: ModelInfo,
     pub(crate) tool_mode: ToolMode,
+    pub(crate) multi_agent_version: Option<MultiAgentVersion>,
     pub(crate) session_telemetry: SessionTelemetry,
     pub(crate) provider: SharedModelProvider,
     pub(crate) reasoning_effort: Option<ReasoningEffortConfig>,
@@ -184,6 +187,7 @@ impl TurnContext {
             .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 truncation_policy = model_info.truncation_policy.into();
         let supported_reasoning_levels = model_info
             .supported_reasoning_levels
@@ -225,6 +229,7 @@ impl TurnContext {
             auth_manager: self.auth_manager.clone(),
             model_info: model_info.clone(),
             tool_mode,
+            multi_agent_version,
             session_telemetry: self
                 .session_telemetry
                 .clone()
@@ -489,6 +494,8 @@ 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);
         per_turn_config.service_tier = get_service_tier(
             per_turn_config.service_tier,
             per_turn_config.features.enabled(Feature::FastMode),
@@ -516,6 +523,7 @@ impl Session {
             auth_manager: auth_manager_for_context,
             model_info: model_info.clone(),
             tool_mode,
+            multi_agent_version,
             session_telemetry: session_telemetry_for_context,
             provider: provider_for_context,
             reasoning_effort,

codex-rs/core/src/tools/spec_plan.rs

@@ -292,23 +292,11 @@ fn code_mode_only_enabled(turn_context: &TurnContext) -> bool {
 }
 
 fn multi_agent_v2_enabled(turn_context: &TurnContext) -> bool {
-    multi_agent_version(turn_context) == Some(MultiAgentVersion::V2)
+    turn_context.multi_agent_version == Some(MultiAgentVersion::V2)
 }
 
 fn collab_tools_enabled(turn_context: &TurnContext) -> bool {
-    multi_agent_version(turn_context).is_some()
-}
-
-fn multi_agent_version(turn_context: &TurnContext) -> Option<MultiAgentVersion> {
-    turn_context.model_info.multi_agent_version.or_else(|| {
-        if turn_context.features.get().enabled(Feature::MultiAgentV2) {
-            Some(MultiAgentVersion::V2)
-        } else if turn_context.features.get().enabled(Feature::Collab) {
-            Some(MultiAgentVersion::V1)
-        } else {
-            None
-        }
-    })
+    turn_context.multi_agent_version.is_some()
 }
 
 fn goal_tools_enabled(turn_context: &TurnContext) -> bool {

codex-rs/core/src/tools/spec_plan_tests.rs

@@ -31,6 +31,7 @@ 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;
@@ -218,7 +219,7 @@ fn set_feature(turn: &mut TurnContext, feature: Feature, enabled: bool) {
             .expect("test feature should be disableable in config");
     }
     turn.config = Arc::new(config);
-    resolve_tool_mode_for_turn(turn);
+    resolve_model_selectors_for_turn(turn);
 }
 
 fn set_features(turn: &mut TurnContext, features: &[Feature]) {
@@ -233,8 +234,9 @@ fn update_config(turn: &mut TurnContext, update: impl FnOnce(&mut crate::config:
     turn.config = Arc::new(config);
 }
 
-fn resolve_tool_mode_for_turn(turn: &mut TurnContext) {
+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) {
@@ -810,7 +812,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_tool_mode_for_turn(turn);
+        resolve_model_selectors_for_turn(turn);
     })
     .await;
     direct.assert_visible_lacks(&[
@@ -824,6 +826,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);
     })
     .await;
 
@@ -840,6 +843,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);
     })
     .await;