From 0c3f349afaef781ee90050a239de95989dd3f559 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Sun, 31 May 2026 11:35:24 -0700 Subject: [PATCH] codex: resolve multi-agent version during session spawn --- codex-rs/core/src/agent/control.rs | 9 +-- codex-rs/core/src/codex_delegate.rs | 2 +- codex-rs/core/src/session/mod.rs | 14 ++-- codex-rs/core/src/session/session.rs | 28 +++++++ .../core/src/session/tests/guardian_tests.rs | 2 +- codex-rs/core/src/thread_manager.rs | 80 ++++++++----------- 6 files changed, 74 insertions(+), 61 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1e37a07703..14449946fe 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/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, ) diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index e54e166a6e..5e298b4a5f 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/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??; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5f1e69c7eb..e2e9060858 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, - pub(crate) multi_agent_version: Option, + pub(crate) inherited_multi_agent_version: Option, } 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() diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 937b5a08d7..4689574381 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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(); diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 2c64144a3e..8b1cf15016 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/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"); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 12f23d7bc3..83d5195f49 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/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, +) -> Option { + 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, inherited_multi_agent_version: Option, @@ -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?;