From 2e4bda03321433a017b3edfc25c5eb34e34bbb21 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Sun, 31 May 2026 12:31:11 -0700 Subject: [PATCH] codex: simplify multi-agent version resolution --- codex-rs/core/src/agent/control.rs | 83 +++++++++++------------------ codex-rs/core/src/session/mod.rs | 6 ++- codex-rs/core/src/thread_manager.rs | 69 ++++++------------------ 3 files changed, 49 insertions(+), 109 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 14449946fe..b49ca849fd 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -144,54 +144,6 @@ 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_inherited_multi_agent_version( - state: &ThreadManagerState, - thread_id: ThreadId, - session_source: &SessionSource, -) -> CodexResult<( - InitialHistory, - Option, - Option, -)> { - let inherited_multi_agent_version = - if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }) = session_source - { - match state.get_thread(*parent_thread_id).await { - Ok(parent_thread) => parent_thread.multi_agent_version(), - Err(_) => None, - } - } else { - None - }; - let stored_thread = state - .read_stored_thread(ReadThreadParams { - thread_id, - include_archived: true, - include_history: true, - }) - .await?; - let history = stored_thread - .history - .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))? - .items; - let initial_history = InitialHistory::Resumed(ResumedHistory { - conversation_id: thread_id, - history, - rollout_path: stored_thread.rollout_path, - }); - let parent_thread_id = stored_thread.parent_thread_id; - let multi_agent_version = state - .resolve_inherited_multi_agent_version( - &initial_history, - /*forked_from_thread_id*/ None, - inherited_multi_agent_version, - ) - .await; - Ok((initial_history, parent_thread_id, multi_agent_version)) -} - /// Control-plane handle for multi-agent operations. /// `AgentControl` is held by each session (via `SessionServices`). It provides capability to /// spawn new agents and the inter-agent communication layer. @@ -670,13 +622,38 @@ impl AgentControl { ) -> CodexResult { 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_inherited_multi_agent_version( - state.as_ref(), + let inherited_multi_agent_version = + if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) = &session_source + { + match state.get_thread(*parent_thread_id).await { + Ok(parent_thread) => parent_thread.multi_agent_version(), + Err(_) => None, + } + } else { + None + }; + let stored_thread = state + .read_stored_thread(ReadThreadParams { thread_id, - &session_source, - ) + include_archived: true, + include_history: true, + }) .await?; + let history = stored_thread + .history + .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))? + .items; + let initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history, + rollout_path: stored_thread.rollout_path, + }); + let parent_thread_id = stored_thread.parent_thread_id; + let multi_agent_version = initial_history + .get_multi_agent_version() + .or(inherited_multi_agent_version); let agent_max_threads = config .effective_agent_max_threads(multi_agent_version) .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index e2e9060858..4835575918 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -540,8 +540,10 @@ 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 multi_agent_version = conversation_history + .get_multi_agent_version() + .or(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()))?; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 83d5195f49..25cdc77f34 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -96,17 +96,6 @@ 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, } @@ -951,17 +940,9 @@ impl ThreadManager { InitialHistory::Forked(_) => history.forked_from_id(), InitialHistory::New | InitialHistory::Cleared => None, }; - let inherited_multi_agent_version = self - .state - .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, - inherited_multi_agent_version, + history.get_multi_agent_version(), ); let history = fork_history_from_snapshot(snapshot, history, interrupted_marker); let environments = default_thread_environment_selections( @@ -982,7 +963,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, - inherited_multi_agent_version, + /*inherited_multi_agent_version*/ None, )) .await } @@ -1232,29 +1213,6 @@ impl ThreadManagerState { .await } - pub(crate) async fn resolve_inherited_multi_agent_version( - &self, - initial_history: &InitialHistory, - forked_from_thread_id: Option, - inherited_multi_agent_version: Option, - ) -> Option { - if let Some(multi_agent_version) = initial_history.get_multi_agent_version() { - return Some(multi_agent_version); - } - - 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, - }; - let live_multi_agent_version = match source_thread.as_ref() { - Some(source_thread) => source_thread.multi_agent_version(), - None => None, - }; - live_multi_agent_version.or(inherited_multi_agent_version) - } - /// Spawn a new thread with optional history and register it with the manager. #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_thread( @@ -1345,18 +1303,21 @@ impl ThreadManagerState { .parent_rollout_thread_trace_for_source(&session_source, &initial_history) .await; let tracked_session_source = session_source.clone(); - let source_thread_id = - multi_agent_version_source_thread_id(&initial_history, forked_from_thread_id); + 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 inherited_multi_agent_version = match inherited_multi_agent_version { Some(multi_agent_version) => Some(multi_agent_version), - None => { - self.resolve_inherited_multi_agent_version( - &initial_history, - forked_from_thread_id, - /*inherited_multi_agent_version*/ None, - ) - .await - } + None => match source_thread_id { + Some(source_thread_id) => self + .get_thread(source_thread_id) + .await + .ok() + .and_then(|source_thread| source_thread.multi_agent_version()), + None => None, + }, }; let CodexSpawnOk { codex, thread_id, ..