From d039da1a575d8f3bcbcc8fe9cdf190fe7c15e63c Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Sat, 30 May 2026 21:25:56 -0700 Subject: [PATCH] codex: simplify multi-agent inheritance --- codex-rs/app-server/tests/common/rollout.rs | 5 +- codex-rs/core/src/agent/control.rs | 46 +++++++++++-------- codex-rs/core/src/agent/control_tests.rs | 8 +++- codex-rs/core/src/session/review.rs | 2 +- codex-rs/core/src/thread_manager.rs | 27 ++++++----- .../core/src/tools/handlers/agent_jobs.rs | 2 +- .../src/tools/handlers/multi_agents/spawn.rs | 2 +- .../src/tools/handlers/multi_agents_tests.rs | 6 +-- .../tools/handlers/multi_agents_v2/spawn.rs | 2 +- 9 files changed, 60 insertions(+), 40 deletions(-) diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 6d446e1c42..408ef646fd 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -2,6 +2,7 @@ use anyhow::Result; use codex_protocol::ThreadId; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GitInfo; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; @@ -194,7 +195,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( base_instructions: None, dynamic_tools: None, memory_mode: None, - multi_agent_version: None, + multi_agent_version: Some(MultiAgentVersion::V1), }; let payload = serde_json::to_value(SessionMetaLine { meta, @@ -280,7 +281,7 @@ pub fn create_fake_rollout_with_text_elements( base_instructions: None, dynamic_tools: None, memory_mode: None, - multi_agent_version: None, + multi_agent_version: Some(MultiAgentVersion::V1), }; let payload = serde_json::to_value(SessionMetaLine { meta, diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index e21326a6dc..eb3a0a61d7 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -57,7 +57,7 @@ pub(crate) struct SpawnAgentOptions { pub(crate) fork_mode: Option, pub(crate) parent_thread_id: Option, pub(crate) environments: Option>, - pub(crate) multi_agent_version: Option>, + pub(crate) multi_agent_version: Option, } #[derive(Clone, Debug)] @@ -191,11 +191,18 @@ impl AgentControl { initial_operation: Op, session_source: Option, ) -> CodexResult { + let multi_agent_version = session_source + .is_none() + .then(|| config.multi_agent_version_from_features()) + .flatten(); let spawned_agent = Box::pin(self.spawn_agent_internal( config, initial_operation, session_source, - SpawnAgentOptions::default(), + SpawnAgentOptions { + multi_agent_version, + ..Default::default() + }, )) .await?; Ok(spawned_agent.thread_id) @@ -221,15 +228,18 @@ impl AgentControl { options: SpawnAgentOptions, ) -> CodexResult { let state = self.upgrade()?; - let inherited_multi_agent_version = self + let multi_agent_version = self .inherited_multi_agent_version_for_source( &state, session_source.as_ref(), options.multi_agent_version, ) - .await; + .await + .ok_or_else(|| { + CodexErr::InvalidRequest("multi-agent version is unresolved".to_string()) + })?; let agent_max_threads = config - .effective_agent_max_threads(inherited_multi_agent_version.flatten()) + .effective_agent_max_threads(Some(multi_agent_version)) .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; let inherited_shell_snapshot = self @@ -271,7 +281,7 @@ impl AgentControl { &options, inherited_shell_snapshot, inherited_exec_policy, - inherited_multi_agent_version, + Some(multi_agent_version), )) .await? } @@ -288,7 +298,7 @@ impl AgentControl { inherited_shell_snapshot, inherited_exec_policy, options.environments.clone(), - inherited_multi_agent_version, + Some(multi_agent_version), )) .await? } @@ -385,7 +395,7 @@ impl AgentControl { options: &SpawnAgentOptions, inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { if options.fork_parent_spawn_call_id.is_none() { return Err(CodexErr::Fatal( @@ -529,7 +539,7 @@ impl AgentControl { inherited_shell_snapshot, inherited_exec_policy, options.environments.clone(), - Some(multi_agent_version), + multi_agent_version, ) .await } @@ -700,7 +710,7 @@ impl AgentControl { parent_thread_id, inherited_shell_snapshot, inherited_exec_policy, - inherited_multi_agent_version: Some(multi_agent_version), + inherited_multi_agent_version: multi_agent_version, }) .await?; let mut agent_metadata = agent_metadata; @@ -1198,19 +1208,19 @@ impl AgentControl { &self, state: &Arc, session_source: Option<&SessionSource>, - requested_multi_agent_version: Option>, - ) -> Option> { + requested_multi_agent_version: Option, + ) -> Option { if requested_multi_agent_version.is_some() { return requested_multi_agent_version; } - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + if let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, .. })) = session_source - else { - return None; - }; - let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; - Some(parent_thread.multi_agent_version().await) + && let Ok(parent_thread) = state.get_thread(*parent_thread_id).await + { + return parent_thread.multi_agent_version().await; + } + None } async fn inherited_exec_policy_for_source( diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 9b67c75a8c..11f23c9350 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1638,9 +1638,15 @@ async fn spawn_child_completion_notifies_parent_history() { #[tokio::test] async fn multi_agent_v2_completion_ignores_dead_direct_parent() { let harness = AgentControlHarness::new().await; - let (root_thread_id, root_thread) = harness.start_thread().await; let mut config = harness.config.clone(); let _ = config.features.enable(Feature::MultiAgentV2); + let root = harness + .manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + let root_thread_id = root.thread_id; + let root_thread = root.thread; let worker_path = AgentPath::root().join("worker_a").expect("worker path"); let worker_thread_id = harness .control diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 9dff6fc051..d2804a039a 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -128,7 +128,7 @@ pub(super) async fn spawn_review_thread( user_instructions: None, compact_prompt: parent_turn_context.compact_prompt.clone(), collaboration_mode: parent_turn_context.collaboration_mode.clone(), - multi_agent_version: parent_turn_context.multi_agent_version, + multi_agent_version: None, personality: parent_turn_context.personality, approval_policy: parent_turn_context.approval_policy.clone(), permission_profile: parent_turn_context.permission_profile(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 82b05cd58a..1a56ab5804 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -194,7 +194,7 @@ pub(crate) struct ResumeThreadWithHistoryOptions { pub(crate) parent_thread_id: Option, pub(crate) inherited_shell_snapshot: Option>, pub(crate) inherited_exec_policy: Option>, - pub(crate) inherited_multi_agent_version: Option>, + pub(crate) inherited_multi_agent_version: Option, } /// Shared, `Arc`-owned state for [`ThreadManager`]. This `Arc` is required to have a single @@ -605,7 +605,7 @@ impl ThreadManager { &self, options: StartThreadOptions, forked_from_thread_id: Option, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let (resumed_session_source, resumed_thread_source) = options .initial_history @@ -657,13 +657,16 @@ impl ThreadManager { )) })?; let history = stored_thread_to_initial_history(stored_thread, fork_source.rollout_path())?; - let inherited_multi_agent_version = fork_source.multi_agent_version().await; + let inherited_multi_agent_version = + fork_source.multi_agent_version().await.ok_or_else(|| { + CodexErr::InvalidRequest("multi-agent version is unresolved".to_string()) + })?; options.initial_history = fork_history_from_snapshot( ForkSnapshot::Interrupted, history, InterruptedTurnHistoryMarker::from_config_and_version( &options.config, - inherited_multi_agent_version, + Some(inherited_multi_agent_version), ), ); self.start_thread_with_options_and_fork_source( @@ -969,7 +972,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, - Some(multi_agent_version), + multi_agent_version, )) .await } @@ -1112,7 +1115,7 @@ impl ThreadManagerState { inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, environments: Option>, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let environments = environments.unwrap_or_else(|| { default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) @@ -1192,7 +1195,7 @@ impl ThreadManagerState { inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, environments: Option>, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let environments = environments.unwrap_or_else(|| { default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) @@ -1224,14 +1227,14 @@ impl ThreadManagerState { config: &Config, initial_history: &InitialHistory, forked_from_thread_id: Option, - inherited_multi_agent_version: 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 multi_agent_version = inherited_multi_agent_version - .unwrap_or_else(|| config.multi_agent_version_from_features()); + let multi_agent_version = + inherited_multi_agent_version.or_else(|| config.multi_agent_version_from_features()); let multi_agent_version = multi_agent_version?; let source_thread_id = match initial_history { InitialHistory::Resumed(resumed) => Some(resumed.conversation_id), @@ -1277,7 +1280,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, user_shell_override: Option, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { Box::pin(self.spawn_thread_with_source( config, @@ -1320,7 +1323,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, user_shell_override: Option, - inherited_multi_agent_version: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 1ba8de62fb..111c10273e 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -216,7 +216,7 @@ async fn run_agent_job_loop( SpawnAgentOptions { parent_thread_id: Some(session.conversation_id), environments: Some(turn.environments.to_selections()), - multi_agent_version: Some(turn.multi_agent_version), + multi_agent_version: turn.multi_agent_version, ..Default::default() }, ) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 1c90ac7dfe..15990b4529 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -125,7 +125,7 @@ async fn handle_spawn_agent( fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), parent_thread_id: Some(session.conversation_id), environments: Some(turn.environments.to_selections()), - multi_agent_version: Some(turn.multi_agent_version), + multi_agent_version: turn.multi_agent_version, }, )) .await diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 9d8250a6cf..7cc9dbdf04 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -1881,15 +1881,15 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); + let mut config = turn.config.as_ref().clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + set_turn_config(&mut turn, config); let root = manager .start_thread((*turn.config).clone()) .await .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - let mut config = turn.config.as_ref().clone(); - let _ = config.features.enable(Feature::MultiAgentV2); - set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index a5309bdef0..d7d6e9f555 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -145,7 +145,7 @@ async fn handle_spawn_agent( fork_mode, parent_thread_id: Some(session.conversation_id), environments: Some(turn.environments.to_selections()), - multi_agent_version: Some(turn.multi_agent_version), + multi_agent_version: turn.multi_agent_version, }, ), )