codex: resolve multi-agent version during session spawn

codex-rs/core/src/agent/control.rs

@@ -144,9 +144,8 @@ fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &
         .any(|usage_hint_text| usage_hint_text == text)
 }
 
-async fn load_resumed_history_and_resolve_multi_agent_version(
+async fn load_resumed_history_and_resolve_inherited_multi_agent_version(
     state: &ThreadManagerState,
-    config: &Config,
     thread_id: ThreadId,
     session_source: &SessionSource,
 ) -> CodexResult<(
@@ -184,8 +183,7 @@ async fn load_resumed_history_and_resolve_multi_agent_version(
     });
     let parent_thread_id = stored_thread.parent_thread_id;
     let multi_agent_version = state
-        .resolve_multi_agent_version(
-            config,
+        .resolve_inherited_multi_agent_version(
             &initial_history,
             /*forked_from_thread_id*/ None,
             inherited_multi_agent_version,
@@ -673,9 +671,8 @@ impl AgentControl {
         let state = self.upgrade()?;
         let state_db_ctx = state.state_db();
         let (initial_history, parent_thread_id, multi_agent_version) =
-            load_resumed_history_and_resolve_multi_agent_version(
+            load_resumed_history_and_resolve_inherited_multi_agent_version(
                 state.as_ref(),
-                &config,
                 thread_id,
                 &session_source,
             )

codex-rs/core/src/codex_delegate.rs

@@ -106,7 +106,7 @@ pub(crate) async fn run_codex_thread_interactive(
         analytics_events_client: Some(parent_session.services.analytics_events_client.clone()),
         thread_store: Arc::clone(&parent_session.services.thread_store),
         attestation_provider: parent_session.services.attestation_provider.clone(),
-        multi_agent_version,
+        inherited_multi_agent_version: multi_agent_version,
     }))
     .or_cancel(&cancel_token)
     .await??;

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

@@ -138,6 +138,7 @@ use codex_thread_store::LiveThreadInitGuard;
 use codex_thread_store::LocalThreadStore;
 use codex_thread_store::ReadThreadParams;
 use codex_thread_store::ResumeThreadParams;
+use codex_thread_store::SetMultiAgentVersionIfUnsetParams;
 use codex_thread_store::ThreadEventPersistenceMode;
 use codex_thread_store::ThreadPersistenceMetadata;
 use codex_thread_store::ThreadStore;
@@ -421,7 +422,7 @@ pub(crate) struct CodexSpawnArgs {
     pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
     pub(crate) thread_store: Arc<dyn ThreadStore>,
     pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
-    pub(crate) multi_agent_version: Option<MultiAgentVersion>,
+    pub(crate) inherited_multi_agent_version: Option<MultiAgentVersion>,
 }
 
 pub(crate) const INITIAL_SUBMIT_ID: &str = "";
@@ -484,15 +485,11 @@ impl Codex {
             analytics_events_client,
             thread_store,
             attestation_provider,
-            multi_agent_version,
+            inherited_multi_agent_version,
         } = args;
         let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
         let (tx_event, rx_event) = async_channel::unbounded();
 
-        let _ = config
-            .effective_agent_max_threads(multi_agent_version)
-            .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?;
-
         let primary_environment = environment_selections.primary_environment();
         let mut user_instruction_warnings = Vec::new();
         let user_instructions = AgentsMdManager::new(&config)
@@ -543,6 +540,11 @@ impl Codex {
         let model_info = models_manager
             .get_model_info(model.as_str(), &config.to_models_manager_config())
             .await;
+        let multi_agent_version =
+            inherited_multi_agent_version.or_else(|| config.multi_agent_version_from_features());
+        let _ = config
+            .effective_agent_max_threads(multi_agent_version)
+            .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?;
         let base_instructions = config
             .base_instructions
             .clone()

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

@@ -522,6 +522,34 @@ impl Session {
             .parent_thread_id
             .or_else(|| initial_history.get_resumed_parent_thread_id());
         session_configuration.parent_thread_id = parent_thread_id;
+        let source_thread_id = match &initial_history {
+            InitialHistory::Resumed(resumed) => Some(resumed.conversation_id),
+            InitialHistory::Forked(_) => forked_from_id,
+            InitialHistory::New | InitialHistory::Cleared => None,
+        };
+        let multi_agent_version = match (source_thread_id, multi_agent_version) {
+            (Some(source_thread_id), Some(multi_agent_version))
+                if !config.ephemeral && initial_history.get_multi_agent_version().is_none() =>
+            {
+                match thread_store
+                    .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams {
+                        thread_id: source_thread_id,
+                        multi_agent_version,
+                        include_archived: true,
+                    })
+                    .await
+                {
+                    Ok(multi_agent_version) => Some(multi_agent_version),
+                    Err(err) => {
+                        warn!(
+                            "failed to lock multi-agent version for legacy thread {source_thread_id}: {err}"
+                        );
+                        Some(multi_agent_version)
+                    }
+                }
+            }
+            (_, multi_agent_version) => multi_agent_version,
+        };
         let multi_agent_version = multi_agent_version.map(OnceLock::from).unwrap_or_default();
         let initial_multi_agent_version = multi_agent_version.get().copied();
 

codex-rs/core/src/session/tests/guardian_tests.rs

@@ -710,7 +710,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
         analytics_events_client: None,
         thread_store,
         attestation_provider: None,
-        multi_agent_version: None,
+        inherited_multi_agent_version: None,
     })
     .await
     .expect("spawn guardian subagent");

codex-rs/core/src/thread_manager.rs

@@ -58,7 +58,6 @@ use codex_thread_store::LocalThreadStore;
 use codex_thread_store::LocalThreadStoreConfig;
 use codex_thread_store::ReadThreadByRolloutPathParams;
 use codex_thread_store::ReadThreadParams;
-use codex_thread_store::SetMultiAgentVersionIfUnsetParams;
 use codex_thread_store::StoredThread;
 use codex_thread_store::ThreadMetadataPatch;
 use codex_thread_store::ThreadStore;
@@ -97,6 +96,17 @@ fn should_use_test_thread_manager_behavior() -> bool {
     FORCE_TEST_THREAD_MANAGER_BEHAVIOR.load(Ordering::Relaxed)
 }
 
+fn multi_agent_version_source_thread_id(
+    initial_history: &InitialHistory,
+    forked_from_thread_id: Option<ThreadId>,
+) -> Option<ThreadId> {
+    match initial_history {
+        InitialHistory::Resumed(resumed) => Some(resumed.conversation_id),
+        InitialHistory::Forked(_) => forked_from_thread_id,
+        InitialHistory::New | InitialHistory::Cleared => None,
+    }
+}
+
 struct TempCodexHomeGuard {
     path: PathBuf,
 }
@@ -941,17 +951,18 @@ impl ThreadManager {
             InitialHistory::Forked(_) => history.forked_from_id(),
             InitialHistory::New | InitialHistory::Cleared => None,
         };
-        let multi_agent_version = self
+        let inherited_multi_agent_version = self
             .state
-            .resolve_multi_agent_version(
-                &config,
+            .resolve_inherited_multi_agent_version(
                 &history,
                 forked_from_thread_id,
                 /*inherited_multi_agent_version*/ None,
             )
             .await;
-        let interrupted_marker =
-            InterruptedTurnHistoryMarker::from_config_and_version(&config, multi_agent_version);
+        let interrupted_marker = InterruptedTurnHistoryMarker::from_config_and_version(
+            &config,
+            inherited_multi_agent_version,
+        );
         let history = fork_history_from_snapshot(snapshot, history, interrupted_marker);
         let environments = default_thread_environment_selections(
             self.state.environment_manager.as_ref(),
@@ -971,7 +982,7 @@ impl ThreadManager {
             parent_trace,
             environments,
             /*user_shell_override*/ None,
-            multi_agent_version,
+            inherited_multi_agent_version,
         ))
         .await
     }
@@ -1221,9 +1232,8 @@ impl ThreadManagerState {
         .await
     }
 
-    pub(crate) async fn resolve_multi_agent_version(
+    pub(crate) async fn resolve_inherited_multi_agent_version(
         &self,
-        config: &Config,
         initial_history: &InitialHistory,
         forked_from_thread_id: Option<ThreadId>,
         inherited_multi_agent_version: Option<MultiAgentVersion>,
@@ -1232,11 +1242,8 @@ impl ThreadManagerState {
             return Some(multi_agent_version);
         }
 
-        let source_thread_id = match initial_history {
-            InitialHistory::Resumed(resumed) => Some(resumed.conversation_id),
-            InitialHistory::Forked(_) => forked_from_thread_id,
-            InitialHistory::New | InitialHistory::Cleared => None,
-        };
+        let source_thread_id =
+            multi_agent_version_source_thread_id(initial_history, forked_from_thread_id);
         let source_thread = match source_thread_id {
             Some(source_thread_id) => self.get_thread(source_thread_id).await.ok(),
             None => None,
@@ -1245,35 +1252,7 @@ impl ThreadManagerState {
             Some(source_thread) => source_thread.multi_agent_version(),
             None => None,
         };
-        let multi_agent_version = live_multi_agent_version
-            .or(inherited_multi_agent_version)
-            .or_else(|| config.multi_agent_version_from_features())?;
-        let Some(source_thread_id) = source_thread_id else {
-            return Some(multi_agent_version);
-        };
-        let multi_agent_version = match self
-            .thread_store
-            .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams {
-                thread_id: source_thread_id,
-                multi_agent_version,
-                include_archived: true,
-            })
-            .await
-        {
-            Ok(multi_agent_version) => multi_agent_version,
-            Err(err) => {
-                warn!(
-                    "failed to lock multi-agent version for legacy thread {source_thread_id}: {err}"
-                );
-                multi_agent_version
-            }
-        };
-        match source_thread {
-            Some(source_thread) => {
-                Some(source_thread.set_multi_agent_version_if_unset(multi_agent_version))
-            }
-            None => Some(multi_agent_version),
-        }
+        live_multi_agent_version.or(inherited_multi_agent_version)
     }
 
     /// Spawn a new thread with optional history and register it with the manager.
@@ -1366,11 +1345,12 @@ impl ThreadManagerState {
             .parent_rollout_thread_trace_for_source(&session_source, &initial_history)
             .await;
         let tracked_session_source = session_source.clone();
-        let multi_agent_version = match inherited_multi_agent_version {
+        let source_thread_id =
+            multi_agent_version_source_thread_id(&initial_history, forked_from_thread_id);
+        let inherited_multi_agent_version = match inherited_multi_agent_version {
             Some(multi_agent_version) => Some(multi_agent_version),
             None => {
-                self.resolve_multi_agent_version(
-                    &config,
+                self.resolve_inherited_multi_agent_version(
                     &initial_history,
                     forked_from_thread_id,
                     /*inherited_multi_agent_version*/ None,
@@ -1408,9 +1388,15 @@ impl ThreadManagerState {
             analytics_events_client: self.analytics_events_client.clone(),
             thread_store: Arc::clone(&self.thread_store),
             attestation_provider: self.attestation_provider.clone(),
-            multi_agent_version,
+            inherited_multi_agent_version,
         })
         .await?;
+        if let Some(source_thread_id) = source_thread_id
+            && let Some(multi_agent_version) = codex.session.multi_agent_version()
+            && let Ok(source_thread) = self.get_thread(source_thread_id).await
+        {
+            source_thread.set_multi_agent_version_if_unset(multi_agent_version);
+        }
         let new_thread = self
             .finalize_thread_spawn(codex, thread_id, tracked_session_source)
             .await?;