diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 63681382ea..6d446e1c42 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -194,6 +194,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; let payload = serde_json::to_value(SessionMetaLine { meta, @@ -279,6 +280,7 @@ pub fn create_fake_rollout_with_text_elements( base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; let payload = serde_json::to_value(SessionMetaLine { meta, diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index eb534e517f..a71e31e3dd 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -129,6 +129,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index df79aea9b4..9e6be687bb 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1363,6 +1363,7 @@ async fn seed_pathless_store_thread( thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index e0263d77f3..8554aea6ef 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -1805,6 +1805,7 @@ stream_max_retries = 0 base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; std::fs::write( &rollout_path, diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 46ce38f48d..f6a31c750e 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -216,6 +216,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index d240f979e3..cb75b32533 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -12,7 +12,6 @@ use crate::shell_snapshot::ShellSnapshot; use crate::thread_manager::ResumeThreadWithHistoryOptions; use crate::thread_manager::ThreadManagerState; use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns; -use codex_features::Feature; use codex_protocol::AgentPath; use codex_protocol::SessionId; use codex_protocol::ThreadId; @@ -24,6 +23,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::ForkedHistory; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; @@ -57,6 +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, } #[derive(Clone, Debug)] @@ -214,12 +215,24 @@ impl AgentControl { async fn spawn_agent_internal( &self, - config: crate::config::Config, + mut config: crate::config::Config, initial_operation: Op, session_source: Option, options: SpawnAgentOptions, ) -> CodexResult { let state = self.upgrade()?; + let inherited_multi_agent_version = self + .inherited_multi_agent_version_for_source( + &state, + session_source.as_ref(), + options.multi_agent_version, + ) + .await; + if let Some(multi_agent_version) = inherited_multi_agent_version { + config + .apply_multi_agent_version(multi_agent_version) + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; + } let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let inherited_shell_snapshot = self .inherited_shell_snapshot_for_source(&state, session_source.as_ref()) @@ -260,6 +273,7 @@ impl AgentControl { &options, inherited_shell_snapshot, inherited_exec_policy, + inherited_multi_agent_version, )) .await? } @@ -276,6 +290,7 @@ impl AgentControl { inherited_shell_snapshot, inherited_exec_policy, options.environments.clone(), + inherited_multi_agent_version, )) .await? } @@ -342,7 +357,7 @@ impl AgentControl { self.send_input(new_thread.thread_id, initial_operation) .await?; - if !new_thread.thread.enabled(Feature::MultiAgentV2) { + if new_thread.thread.multi_agent_version().await != MultiAgentVersion::V2 { let child_reference = agent_metadata .agent_path .as_ref() @@ -371,6 +386,7 @@ impl AgentControl { options: &SpawnAgentOptions, inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { if options.fork_parent_spawn_call_id.is_none() { return Err(CodexErr::Fatal( @@ -414,6 +430,18 @@ impl AgentControl { )) })?; + let multi_agent_version = if let Some(parent_thread) = parent_thread.as_ref() { + parent_thread.multi_agent_version().await + } else { + state + .resolve_multi_agent_version( + &config, + &InitialHistory::Forked(parent_history.items.clone()), + Some(parent_thread_id), + inherited_multi_agent_version, + ) + .await + }; let mut forked_rollout_items = parent_history.items; if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { forked_rollout_items = @@ -421,7 +449,7 @@ impl AgentControl { } let multi_agent_v2_usage_hint_texts_to_filter: Vec = if let Some(parent_thread) = parent_thread.as_ref() { - if parent_thread.enabled(Feature::MultiAgentV2) { + if multi_agent_version == MultiAgentVersion::V2 { let parent_config = parent_thread.codex.session.get_config().await; [ parent_config @@ -439,7 +467,7 @@ impl AgentControl { } else { Vec::new() } - } else if config.features.enabled(Feature::MultiAgentV2) { + } else if multi_agent_version == MultiAgentVersion::V2 { [ config.multi_agent_v2.root_agent_usage_hint_text.clone(), config.multi_agent_v2.subagent_usage_hint_text.clone(), @@ -475,8 +503,7 @@ impl AgentControl { } } if preserve_reference_context_item - && config.features.enabled(Feature::MultiAgentV2) - && config.multi_agent_v2.usage_hint_enabled + && multi_agent_version == MultiAgentVersion::V2 && let Some(subagent_usage_hint_text) = config.multi_agent_v2.subagent_usage_hint_text.clone() && let Some(subagent_usage_hint_message) = @@ -503,6 +530,7 @@ impl AgentControl { inherited_shell_snapshot, inherited_exec_policy, options.environments.clone(), + Some(multi_agent_version), ) .await } @@ -589,15 +617,39 @@ impl AgentControl { thread_id: ThreadId, session_source: SessionSource, ) -> CodexResult { - if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) = &session_source - && *depth >= config.agent_max_depth - && !config.features.enabled(Feature::MultiAgentV2) - { - let _ = config.features.disable(Feature::SpawnCsv); - let _ = config.features.disable(Feature::Collab); - } let state = self.upgrade()?; let state_db_ctx = state.state_db(); + let inherited_multi_agent_version = self + .inherited_multi_agent_version_for_source(&state, Some(&session_source), None) + .await; + 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_multi_agent_version( + &config, + &initial_history, + /*forked_from_thread_id*/ None, + inherited_multi_agent_version, + ) + .await; + config + .apply_multi_agent_version(multi_agent_version) + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -635,32 +687,17 @@ impl AgentControl { let inherited_exec_policy = self .inherited_exec_policy_for_source(&state, Some(&session_source), &config) .await; - 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 parent_thread_id = stored_thread.parent_thread_id; let resumed_thread = state .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { config: config.clone(), - initial_history: InitialHistory::Resumed(ResumedHistory { - conversation_id: thread_id, - history, - rollout_path: stored_thread.rollout_path, - }), + initial_history, agent_control: self.clone(), session_source, parent_thread_id, inherited_shell_snapshot, inherited_exec_policy, + inherited_multi_agent_version: Some(multi_agent_version), }) .await?; let mut agent_metadata = agent_metadata; @@ -669,7 +706,7 @@ impl AgentControl { // Resumed threads are re-registered in-memory and need the same listener // attachment path as freshly spawned threads. state.notify_thread_created(resumed_thread.thread_id); - if !resumed_thread.thread.enabled(Feature::MultiAgentV2) { + if resumed_thread.thread.multi_agent_version().await != MultiAgentVersion::V2 { let child_reference = agent_metadata .agent_path .as_ref() @@ -1054,12 +1091,13 @@ impl AgentControl { }; let child_thread = state.get_thread(child_thread_id).await.ok(); let message = format_subagent_notification_message(child_reference.as_str(), &status); - if child_agent_path.is_some() - && child_thread - .as_ref() - .map(|thread| thread.enabled(Feature::MultiAgentV2)) - .unwrap_or(true) - { + let child_uses_multi_agent_v2 = match child_thread.as_ref() { + Some(child_thread) => { + child_thread.multi_agent_version().await == MultiAgentVersion::V2 + } + None => true, + }; + if child_agent_path.is_some() && child_uses_multi_agent_v2 { let Some(child_agent_path) = child_agent_path.clone() else { return; }; @@ -1153,6 +1191,25 @@ impl AgentControl { parent_thread.codex.session.user_shell().shell_snapshot() } + async fn inherited_multi_agent_version_for_source( + &self, + state: &Arc, + session_source: Option<&SessionSource>, + requested_multi_agent_version: Option, + ) -> Option { + if requested_multi_agent_version.is_some() { + return requested_multi_agent_version; + } + 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) + } + async fn inherited_exec_policy_for_source( &self, state: &Arc, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3e0242c037..006ca5f63e 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -21,6 +21,7 @@ use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; @@ -566,6 +567,53 @@ async fn spawn_agent_creates_thread_and_sends_prompt() { assert_eq!(captured, Some(expected)); } +#[tokio::test] +async fn spawn_agent_inherits_parent_multi_agent_version_over_child_feature_flags() { + let harness = AgentControlHarness::new().await; + let mut parent_config = harness.config.clone(); + parent_config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + let parent = harness + .manager + .start_thread(parent_config) + .await + .expect("start parent thread"); + let mut child_config = harness.config.clone(); + child_config + .features + .disable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + child_config + .features + .disable(Feature::Collab) + .expect("test config should allow feature update"); + + let child_thread_id = harness + .control + .spawn_agent( + child_config, + text_input("spawned"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ) + .await + .expect("spawn_agent should succeed"); + let child = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + + assert_eq!(child.multi_agent_version().await, MultiAgentVersion::V2); +} + #[tokio::test] async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 44a8c0ac6b..7fcbd2b015 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -52,6 +52,7 @@ use codex_login::AuthManager; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::error::CodexErr; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::MultiAgentVersion; #[cfg(test)] use crate::session::completed_session_loop_termination; @@ -71,6 +72,7 @@ pub(crate) async fn run_codex_thread_interactive( cancel_token: CancellationToken, subagent_source: SubAgentSource, initial_history: Option, + multi_agent_version: MultiAgentVersion, ) -> Result { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_ops, rx_ops) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -104,6 +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, })) .or_cancel(&cancel_token) .await??; @@ -175,6 +178,7 @@ pub(crate) async fn run_codex_thread_one_shot( subagent_source: SubAgentSource, final_output_json_schema: Option, initial_history: Option, + multi_agent_version: MultiAgentVersion, ) -> Result { // Use a child token so we can stop the delegate after completion without // requiring the caller to cancel the parent token. @@ -188,6 +192,7 @@ pub(crate) async fn run_codex_thread_one_shot( child_cancel.clone(), subagent_source, initial_history, + multi_agent_version, )) .await?; diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 04f6ae3965..0c7981f064 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -172,6 +172,7 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() { cancel_token, SubAgentSource::Review, /*initial_history*/ None, + codex_protocol::protocol::MultiAgentVersion::None, ), ) .await diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 7a6bc69c5c..8d80cfd530 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -24,6 +24,7 @@ use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AdditionalContextEntry; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; @@ -502,6 +503,10 @@ impl CodexThread { self.codex.session.get_config().await } + pub async fn multi_agent_version(&self) -> MultiAgentVersion { + self.codex.session.multi_agent_version().await + } + /// Refresh the thread's layer-backed user config state from a caller-supplied /// config snapshot. Thread-scoped layers and session-static settings remain /// unchanged. diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 70719843ba..85f594580b 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -9899,6 +9899,43 @@ subagent_usage_hint_text = "" Ok(()) } +#[tokio::test] +async fn multi_agent_version_resolution_prefers_v2_then_v1_then_disabled() -> std::io::Result<()> { + let mut config = test_config().await; + config + .features + .disable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config + .features + .disable(Feature::Collab) + .expect("test config should allow feature update"); + assert_eq!( + config.multi_agent_version_from_features(), + codex_protocol::protocol::MultiAgentVersion::None + ); + + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + assert_eq!( + config.multi_agent_version_from_features(), + codex_protocol::protocol::MultiAgentVersion::V1 + ); + + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + assert_eq!( + config.multi_agent_version_from_features(), + codex_protocol::protocol::MultiAgentVersion::V2 + ); + + Ok(()) +} + #[tokio::test] async fn multi_agent_v2_rejects_agents_max_threads() -> std::io::Result<()> { let codex_home = TempDir::new()?; @@ -9912,11 +9949,13 @@ max_threads = 3 "#, )?; - let err = ConfigBuilder::without_managed_config_for_tests() + let mut config = ConfigBuilder::without_managed_config_for_tests() .codex_home(codex_home.path().to_path_buf()) .fallback_cwd(Some(codex_home.path().to_path_buf())) .build() - .await + .await?; + let err = config + .apply_multi_agent_version(codex_protocol::protocol::MultiAgentVersion::V2) .expect_err("agents.max_threads should conflict with multi_agent_v2"); assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ff666b51e9..46a3d3c1e6 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -99,6 +99,7 @@ use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBufGuard; @@ -825,6 +826,8 @@ pub struct Config { /// Maximum number of agent threads that can be open concurrently. pub agent_max_threads: Option, + /// User-configured legacy thread limit before applying the locked multi-agent runtime. + pub(crate) configured_agent_max_threads: Option, /// Maximum runtime in seconds for agent job workers before they are failed. pub agent_job_max_runtime_seconds: Option, @@ -1281,6 +1284,41 @@ impl ConfigBuilder { } impl Config { + pub(crate) fn multi_agent_version_from_features(&self) -> MultiAgentVersion { + if self.features.enabled(Feature::MultiAgentV2) { + MultiAgentVersion::V2 + } else if self.features.enabled(Feature::Collab) { + MultiAgentVersion::V1 + } else { + MultiAgentVersion::None + } + } + + pub(crate) fn apply_multi_agent_version( + &mut self, + multi_agent_version: MultiAgentVersion, + ) -> std::io::Result<()> { + self.agent_max_threads = match multi_agent_version { + MultiAgentVersion::V2 => { + if self.configured_agent_max_threads.is_some() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "agents.max_threads cannot be set when multi_agent_v2 is enabled", + )); + } + Some( + self.multi_agent_v2 + .max_concurrent_threads_per_session + .saturating_sub(1), + ) + } + MultiAgentVersion::None | MultiAgentVersion::V1 => self + .configured_agent_max_threads + .or(DEFAULT_AGENT_MAX_THREADS), + }; + Ok(()) + } + pub fn legacy_sandbox_policy(&self) -> SandboxPolicy { self.permissions.legacy_sandbox_policy(self.cwd.as_path()) } @@ -3063,28 +3101,22 @@ impl Config { )); } validate_multi_agent_v2_tool_namespace(multi_agent_v2.tool_namespace.as_deref())?; - let agent_max_threads_from_config = cfg.agents.as_ref().and_then(|agents| agents.max_threads); + let configured_agent_max_threads = + cfg.agents.as_ref().and_then(|agents| agents.max_threads); + if configured_agent_max_threads == Some(0) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "agents.max_threads must be at least 1", + )); + } let agent_max_threads = if features.enabled(Feature::MultiAgentV2) { - if agent_max_threads_from_config.is_some() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "agents.max_threads cannot be set when multi_agent_v2 is enabled", - )); - } Some( multi_agent_v2 .max_concurrent_threads_per_session .saturating_sub(1), ) } else { - let agent_max_threads = agent_max_threads_from_config.or(DEFAULT_AGENT_MAX_THREADS); - if agent_max_threads == Some(0) { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "agents.max_threads must be at least 1", - )); - } - agent_max_threads + configured_agent_max_threads.or(DEFAULT_AGENT_MAX_THREADS) }; let agent_max_depth = cfg .agents @@ -3457,6 +3489,7 @@ impl Config { .collect(), tool_output_token_limit: cfg.tool_output_token_limit, agent_max_threads, + configured_agent_max_threads, agent_max_depth, agent_roles, memories: memories_config, diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index aa19628d14..2394affe81 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -19,6 +19,7 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ForkedHistory; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SandboxPolicy; @@ -610,6 +611,7 @@ async fn spawn_guardian_review_session( cancel_token.clone(), SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_string()), initial_history, + MultiAgentVersion::None, )) .await?; diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index ad42a41baa..4cf8ab9d8f 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -59,6 +59,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index 85815f8533..9cea171cf1 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -8,6 +8,7 @@ use codex_features::FeatureToml; use codex_features::FeaturesToml; use codex_features::MultiAgentV2ConfigToml; use codex_protocol::ThreadId; +use codex_protocol::protocol::MultiAgentVersion; use crate::config::Config; use crate::config_lock::ConfigLockReplayOptions; @@ -94,7 +95,7 @@ fn session_configuration_to_lock_config_toml( save_session_resolved_fields(sc, &mut lock_config); } - save_config_resolved_fields(config, &mut lock_config)?; + save_config_resolved_fields(sc, config, &mut lock_config)?; drop_lockfile_inputs(&mut lock_config); Ok(lock_config) @@ -124,6 +125,7 @@ fn save_session_resolved_fields(sc: &SessionConfiguration, lock_config: &mut Con /// Persist the resolved representation so replay compares against the behavior /// Codex actually ran with, not only the user-authored TOML inputs. fn save_config_resolved_fields( + sc: &SessionConfiguration, config: &Config, lock_config: &mut ConfigToml, ) -> anyhow::Result<()> { @@ -144,10 +146,19 @@ fn save_config_resolved_fields( let features = lock_config .features .get_or_insert_with(FeaturesToml::default); - features.materialize_resolved_enabled(config.features.get()); + let mut materialized_features = config.features.get().clone(); + materialized_features.set_enabled( + Feature::Collab, + sc.multi_agent_version != MultiAgentVersion::None, + ); + materialized_features.set_enabled( + Feature::MultiAgentV2, + sc.multi_agent_version == MultiAgentVersion::V2, + ); + features.materialize_resolved_enabled(&materialized_features); let mut multi_agent_v2: MultiAgentV2ConfigToml = resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?; - multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2)); + multi_agent_v2.enabled = Some(sc.multi_agent_version == MultiAgentVersion::V2); features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2)); features.apps_mcp_path_override = Some(FeatureToml::Config(AppsMcpPathOverrideConfigToml { enabled: Some(config.features.enabled(Feature::AppsMcpPathOverride)), @@ -161,7 +172,7 @@ fn save_config_resolved_fields( let agents = lock_config.agents.get_or_insert_with(Default::default); // Multi-agent v2 owns thread fanout through its feature config. Preserve // the legacy agents.max_threads setting only when v2 is disabled. - agents.max_threads = if config.features.enabled(Feature::MultiAgentV2) { + agents.max_threads = if sc.multi_agent_version == MultiAgentVersion::V2 { None } else { config.agent_max_threads @@ -273,6 +284,45 @@ mod tests { assert_eq!(lockfile.version, crate::config_lock::CONFIG_LOCK_VERSION); } + #[tokio::test] + async fn lock_projects_multi_agent_version_over_current_feature_flags() { + let mut sc = crate::session::tests::make_session_configuration_for_tests().await; + let mut config = (*sc.original_config_do_not_use).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + sc.original_config_do_not_use = Arc::new(config); + + for (multi_agent_version, collab, multi_agent_v2) in [ + (MultiAgentVersion::None, false, false), + (MultiAgentVersion::V1, true, false), + (MultiAgentVersion::V2, true, true), + ] { + sc.multi_agent_version = multi_agent_version; + let lockfile = sc.to_config_lockfile_toml().expect("lock should serialize"); + let lock = &lockfile.config; + let entries = lock + .features + .as_ref() + .expect("lock should materialize features") + .entries(); + assert_eq!(entries.get(Feature::Collab.key()), Some(&collab)); + assert_eq!( + entries.get(Feature::MultiAgentV2.key()), + Some(&multi_agent_v2) + ); + assert_eq!( + lock.agents + .as_ref() + .expect("lock should materialize agent config") + .max_threads + .is_none(), + multi_agent_version == MultiAgentVersion::V2 + ); + } + } + #[tokio::test] async fn lock_skips_session_values_when_model_catalog_fields_are_not_saved() { let mut sc = crate::session::tests::make_session_configuration_for_tests().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 565f378f41..856be547b3 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -105,6 +105,7 @@ use codex_protocol::protocol::HasLegacyEvent; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ItemStartedEvent; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RawResponseItemEvent; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::RolloutItem; @@ -420,6 +421,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: MultiAgentVersion, } pub(crate) const INITIAL_SUBMIT_ID: &str = ""; @@ -482,17 +484,14 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + multi_agent_version, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); - if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) = session_source - && depth >= config.agent_max_depth - && !config.features.enabled(Feature::MultiAgentV2) - { - let _ = config.features.disable(Feature::SpawnCsv); - let _ = config.features.disable(Feature::Collab); - } + config + .apply_multi_agent_version(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(); @@ -574,6 +573,7 @@ impl Codex { let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version, model_reasoning_summary: config.model_reasoning_summary, service_tier, developer_instructions: config.developer_instructions.clone(), @@ -1652,7 +1652,7 @@ impl Session { turn_context: &TurnContext, msg: &EventMsg, ) { - if !self.enabled(Feature::MultiAgentV2) { + if turn_context.multi_agent_version != MultiAgentVersion::V2 { return; } @@ -2623,6 +2623,11 @@ impl Session { state.session_configuration.collaboration_mode.clone() } + pub(crate) async fn multi_agent_version(&self) -> MultiAgentVersion { + let state = self.state.lock().await; + state.session_configuration.multi_agent_version + } + async fn send_raw_response_items(&self, turn_context: &TurnContext, items: &[ResponseItem]) { for item in items { self.send_event( diff --git a/codex-rs/core/src/session/multi_agents.rs b/codex-rs/core/src/session/multi_agents.rs index eb3da0da4b..a78a9c938e 100644 --- a/codex-rs/core/src/session/multi_agents.rs +++ b/codex-rs/core/src/session/multi_agents.rs @@ -1,5 +1,5 @@ use crate::session::turn_context::TurnContext; -use codex_features::Feature; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; @@ -7,7 +7,7 @@ pub(super) fn usage_hint_text<'a>( turn_context: &'a TurnContext, session_source: &SessionSource, ) -> Option<&'a str> { - if !turn_context.features.enabled(Feature::MultiAgentV2) { + if turn_context.multi_agent_version != MultiAgentVersion::V2 { return None; } diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index c476436560..9dff6fc051 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -128,6 +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, 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/session/session.rs b/codex-rs/core/src/session/session.rs index 995f07d701..6527092272 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -9,6 +9,7 @@ use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSpecialPath; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnEnvironmentSelection; use tokio::sync::Semaphore; @@ -45,6 +46,7 @@ pub(crate) struct SessionConfiguration { pub(super) provider: ModelProviderInfo, pub(super) collaboration_mode: CollaborationMode, + pub(super) multi_agent_version: MultiAgentVersion, pub(super) model_reasoning_summary: Option, pub(super) service_tier: Option, @@ -564,6 +566,7 @@ impl Session { text: session_configuration.base_instructions.clone(), }, dynamic_tools: session_configuration.dynamic_tools.clone(), + multi_agent_version: session_configuration.multi_agent_version, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c392825652..43dfb20107 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3079,6 +3079,7 @@ async fn set_rate_limits_retains_previous_credits() { let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -3185,6 +3186,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -3464,6 +3466,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -3715,6 +3718,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -4460,6 +4464,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -4571,6 +4576,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -4808,6 +4814,7 @@ async fn make_session_with_config_and_rx( let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -4913,6 +4920,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), @@ -5966,6 +5974,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -6419,6 +6428,7 @@ where let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, + multi_agent_version: config.multi_agent_version_from_features(), model_reasoning_summary: config.model_reasoning_summary, developer_instructions: config.developer_instructions.clone(), user_instructions: config.user_instructions.clone(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a0625d65e2..4881211bd8 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -710,6 +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: codex_protocol::protocol::MultiAgentVersion::None, }) .await .expect("spawn guardian subagent"); @@ -729,6 +730,10 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { }], } ); + assert_eq!( + codex.session.multi_agent_version().await, + codex_protocol::protocol::MultiAgentVersion::None + ); drop(codex); } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 1e0a3f17c5..28fb51150d 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -8,6 +8,7 @@ use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile; @@ -78,6 +79,7 @@ pub struct TurnContext { pub(crate) compact_prompt: Option, pub(crate) user_instructions: Option, pub(crate) collaboration_mode: CollaborationMode, + pub(crate) multi_agent_version: MultiAgentVersion, pub(crate) personality: Option, pub(crate) approval_policy: Constrained, pub(crate) permission_profile: PermissionProfile, @@ -543,6 +545,7 @@ impl Session { compact_prompt: session_configuration.compact_prompt.clone(), user_instructions: session_configuration.user_instructions.clone(), collaboration_mode: session_configuration.collaboration_mode.clone(), + multi_agent_version: session_configuration.multi_agent_version, personality: session_configuration.personality, approval_policy: session_configuration.approval_policy.clone(), permission_profile: session_configuration.permission_profile(), diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 38a8a96276..76d5e6478e 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -44,6 +44,7 @@ use codex_otel::TURN_TOKEN_USAGE_METRIC; use codex_otel::TURN_TOOL_CALL_METRIC; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; @@ -70,11 +71,14 @@ pub(crate) enum InterruptedTurnHistoryMarker { } impl InterruptedTurnHistoryMarker { - pub(crate) fn from_config(config: &Config) -> Self { + pub(crate) fn from_config_and_version( + config: &Config, + multi_agent_version: MultiAgentVersion, + ) -> Self { if !config.agent_interrupt_message_enabled { return Self::Disabled; } - if config.features.enabled(Feature::MultiAgentV2) { + if multi_agent_version == MultiAgentVersion::V2 { Self::Developer } else { Self::ContextualUser @@ -841,7 +845,10 @@ impl Session { if reason == TurnAbortReason::Interrupted && let Some(marker) = interrupted_turn_history_marker( - InterruptedTurnHistoryMarker::from_config(task.turn_context.config.as_ref()), + InterruptedTurnHistoryMarker::from_config_and_version( + task.turn_context.config.as_ref(), + task.turn_context.multi_agent_version, + ), ) { self.record_conversation_items( diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index 9f7eba9573..725fdff4cf 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExitedReviewModeEvent; use codex_protocol::protocol::ItemCompletedEvent; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ReviewOutputEvent; use codex_protocol::protocol::SubAgentSource; use tokio_util::sync::CancellationToken; @@ -132,6 +133,7 @@ async fn start_review_conversation( SubAgentSource::Review, /*final_output_json_schema*/ None, /*initial_history*/ None, + MultiAgentVersion::None, ) .await) .ok() diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 068e789bd3..c01dc2fc98 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -39,6 +39,7 @@ use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ForkedHistory; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; @@ -57,6 +58,7 @@ 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; @@ -192,6 +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, } /// Shared, `Arc`-owned state for [`ThreadManager`]. This `Arc` is required to have a single @@ -591,14 +594,18 @@ impl ThreadManager { &self, options: StartThreadOptions, ) -> CodexResult { - self.start_thread_with_options_and_fork_source(options, /*forked_from_thread_id*/ None) - .await + self.start_thread_with_options_and_fork_source( + options, /*forked_from_thread_id*/ None, + /*inherited_multi_agent_version*/ None, + ) + .await } async fn start_thread_with_options_and_fork_source( &self, options: StartThreadOptions, forked_from_thread_id: Option, + inherited_multi_agent_version: Option, ) -> CodexResult { let (resumed_session_source, resumed_thread_source) = options .initial_history @@ -623,6 +630,7 @@ impl ThreadManager { options.parent_trace, options.environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -649,13 +657,21 @@ 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; options.initial_history = fork_history_from_snapshot( ForkSnapshot::Interrupted, history, - InterruptedTurnHistoryMarker::from_config(&options.config), + InterruptedTurnHistoryMarker::from_config_and_version( + &options.config, + inherited_multi_agent_version, + ), ); - self.start_thread_with_options_and_fork_source(options, Some(forked_from_thread_id)) - .await + self.start_thread_with_options_and_fork_source( + options, + Some(forked_from_thread_id), + Some(inherited_multi_agent_version), + ) + .await } pub async fn resume_thread_from_rollout( @@ -708,6 +724,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, + /*inherited_multi_agent_version*/ None, )) .await } @@ -735,6 +752,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*user_shell_override*/ Some(user_shell_override), + /*inherited_multi_agent_version*/ None, )) .await } @@ -771,6 +789,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*user_shell_override*/ Some(user_shell_override), + /*inherited_multi_agent_version*/ None, )) .await } @@ -920,7 +939,17 @@ impl ThreadManager { InitialHistory::Forked(_) => history.forked_from_id(), InitialHistory::New | InitialHistory::Cleared => None, }; - let interrupted_marker = InterruptedTurnHistoryMarker::from_config(&config); + let multi_agent_version = self + .state + .resolve_multi_agent_version( + &config, + &history, + forked_from_thread_id, + /*inherited_multi_agent_version*/ None, + ) + .await; + let interrupted_marker = + InterruptedTurnHistoryMarker::from_config_and_version(&config, 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(), @@ -940,6 +969,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, + Some(multi_agent_version), )) .await } @@ -1063,6 +1093,7 @@ impl ThreadManagerState { /*inherited_shell_snapshot*/ None, /*inherited_exec_policy*/ None, /*environments*/ None, + /*inherited_multi_agent_version*/ None, )) .await } @@ -1081,6 +1112,7 @@ impl ThreadManagerState { inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, environments: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let environments = environments.unwrap_or_else(|| { default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) @@ -1102,6 +1134,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -1118,6 +1151,7 @@ impl ThreadManagerState { parent_thread_id, inherited_shell_snapshot, inherited_exec_policy, + inherited_multi_agent_version, } = options; let environments = default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd); @@ -1139,6 +1173,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -1157,6 +1192,7 @@ impl ThreadManagerState { inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, environments: Option>, + inherited_multi_agent_version: Option, ) -> CodexResult { let environments = environments.unwrap_or_else(|| { default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) @@ -1178,10 +1214,51 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } + pub(crate) async fn resolve_multi_agent_version( + &self, + config: &Config, + initial_history: &InitialHistory, + forked_from_thread_id: Option, + inherited_multi_agent_version: Option, + ) -> MultiAgentVersion { + if let Some(multi_agent_version) = initial_history.get_multi_agent_version() { + return multi_agent_version; + } + + let multi_agent_version = inherited_multi_agent_version + .unwrap_or_else(|| config.multi_agent_version_from_features()); + 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 Some(source_thread_id) = source_thread_id else { + return 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 + } + } + } + /// Spawn a new thread with optional history and register it with the manager. #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_thread( @@ -1199,6 +1276,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, user_shell_override: Option, + inherited_multi_agent_version: Option, ) -> CodexResult { Box::pin(self.spawn_thread_with_source( config, @@ -1217,6 +1295,7 @@ impl ThreadManagerState { parent_trace, environments, user_shell_override, + inherited_multi_agent_version, )) .await } @@ -1240,6 +1319,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, user_shell_override: Option, + inherited_multi_agent_version: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { @@ -1269,6 +1349,14 @@ impl ThreadManagerState { .parent_rollout_thread_trace_for_source(&session_source, &initial_history) .await; let tracked_session_source = session_source.clone(); + let multi_agent_version = self + .resolve_multi_agent_version( + &config, + &initial_history, + forked_from_thread_id, + inherited_multi_agent_version, + ) + .await; let CodexSpawnOk { codex, thread_id, .. } = Codex::spawn(CodexSpawnArgs { @@ -1299,6 +1387,7 @@ 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, }) .await?; let new_thread = self diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index e768910d51..1beff2421b 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -17,7 +17,10 @@ use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ResumedHistory; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; @@ -765,6 +768,184 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { .expect("shutdown resumed thread"); } +#[tokio::test] +async fn resume_and_fork_preserve_stored_multi_agent_version() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + config + .features + .disable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let manager = ThreadManager::new( + &config, + auth_manager.clone(), + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), + /*analytics_events_client*/ None, + thread_store_from_config(&config, /*state_db*/ None), + /*state_db*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + ); + + let source = manager + .start_thread(config.clone()) + .await + .expect("start source thread"); + assert_eq!( + source.thread.multi_agent_version().await, + MultiAgentVersion::V1 + ); + source.thread.ensure_rollout_materialized().await; + source + .thread + .flush_rollout() + .await + .expect("flush source rollout"); + let rollout_path = source + .thread + .rollout_path() + .expect("source rollout path should exist"); + source + .thread + .shutdown_and_wait() + .await + .expect("shutdown source thread"); + let _ = manager.remove_thread(&source.thread_id).await; + + let mut current_config = config; + current_config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + let resumed = manager + .resume_thread_from_rollout( + current_config.clone(), + rollout_path.clone(), + auth_manager, + /*parent_trace*/ None, + ) + .await + .expect("resume source thread"); + assert_eq!( + resumed.thread.multi_agent_version().await, + MultiAgentVersion::V1 + ); + resumed + .thread + .shutdown_and_wait() + .await + .expect("shutdown resumed thread"); + let _ = manager.remove_thread(&resumed.thread_id).await; + + let forked = manager + .fork_thread( + ForkSnapshot::Interrupted, + current_config, + rollout_path, + /*thread_source*/ None, + /*persist_extended_history*/ false, + /*parent_trace*/ None, + ) + .await + .expect("fork source thread"); + assert_eq!( + forked.thread.multi_agent_version().await, + MultiAgentVersion::V1 + ); + forked + .thread + .shutdown_and_wait() + .await + .expect("shutdown forked thread"); +} + +#[tokio::test] +async fn legacy_resume_continues_when_multi_agent_version_seed_fails() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + config.experimental_thread_store = ThreadStoreConfig::InMemory { + id: format!("thread-manager-{}", uuid::Uuid::new_v4()), + }; + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let thread_store = thread_store_from_config(&config, /*state_db*/ None); + let in_memory_store = thread_store + .as_any() + .downcast_ref::() + .expect("configured in-memory store"); + let manager = ThreadManager::new( + &config, + auth_manager.clone(), + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), + /*analytics_events_client*/ None, + thread_store, + /*state_db*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + ); + let thread_id = ThreadId::new(); + + let resumed = manager + .resume_thread_with_history( + config, + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: vec![RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + cwd: temp_dir.path().to_path_buf(), + ..SessionMeta::default() + }, + git: None, + })], + rollout_path: None, + }), + auth_manager, + /*persist_extended_history*/ false, + /*parent_trace*/ None, + ) + .await + .expect("legacy resume should continue after seed failure"); + assert_eq!( + resumed.thread.multi_agent_version().await, + MultiAgentVersion::V2 + ); + assert_eq!( + in_memory_store + .calls() + .await + .set_multi_agent_version_if_unset, + 1 + ); + resumed + .thread + .shutdown_and_wait() + .await + .expect("shutdown resumed thread"); +} + #[tokio::test] async fn resume_stopped_thread_from_rollout_preserves_thread_source() { let temp_dir = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 821b3619c4..7660b66805 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -213,6 +213,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), ..Default::default() }, ) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs index a536c8d704..0a9d3094a0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs @@ -181,7 +181,7 @@ async fn try_resume_closed_agent( receiver_thread_id: ThreadId, child_depth: i32, ) -> Result<(), FunctionCallError> { - let config = build_agent_resume_config(turn.as_ref(), child_depth)?; + let config = build_agent_resume_config(turn.as_ref())?; Box::pin(session.services.agent_control.resume_agent_from_rollout( config, receiver_thread_id, 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 c4120dd5af..1c90ac7dfe 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -109,7 +109,6 @@ async fn handle_spawn_agent( ) .await?; apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; - apply_spawn_agent_overrides(&mut config, child_depth); let result = Box::pin(session.services.agent_control.spawn_agent_with_metadata( config, @@ -126,6 +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), }, )) .await diff --git a/codex-rs/core/src/tools/handlers/multi_agents_common.rs b/codex-rs/core/src/tools/handlers/multi_agents_common.rs index f1e7bcd429..8a4edbff35 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_common.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_common.rs @@ -8,7 +8,6 @@ use crate::session::turn_context::TurnContext; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; -use codex_features::Feature; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::AgentPath; use codex_protocol::ThreadId; @@ -211,12 +210,8 @@ pub(crate) fn build_agent_spawn_config( Ok(config) } -pub(crate) fn build_agent_resume_config( - turn: &TurnContext, - child_depth: i32, -) -> Result { +pub(crate) fn build_agent_resume_config(turn: &TurnContext) -> Result { let mut config = build_agent_shared_config(turn)?; - apply_spawn_agent_overrides(&mut config, child_depth); // For resume, keep base instructions sourced from rollout/session metadata. config.base_instructions = None; Ok(config) @@ -280,13 +275,6 @@ pub(crate) fn apply_spawn_agent_runtime_overrides( Ok(()) } -pub(crate) fn apply_spawn_agent_overrides(config: &mut Config, child_depth: i32) { - if child_depth >= config.agent_max_depth && !config.features.enabled(Feature::MultiAgentV2) { - let _ = config.features.disable(Feature::SpawnCsv); - let _ = config.features.disable(Feature::Collab); - } -} - pub(crate) async fn apply_requested_spawn_agent_model_overrides( session: &Session, turn: &TurnContext, 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 df8c229411..a9858bc2e7 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -135,6 +135,11 @@ model_reasoning_effort = "minimal" role_name } +fn set_turn_config(turn: &mut TurnContext, config: crate::config::Config) { + turn.multi_agent_version = config.multi_agent_version_from_features(); + turn.config = Arc::new(config); +} + fn expect_text_output(output: T) -> (String, Option) where T: ToolOutput, @@ -383,6 +388,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() { .expect("test config should allow feature update"); let turn = TurnContext { config: Arc::new(config), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::V2, ..turn }; @@ -425,7 +431,7 @@ async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_over .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let err = SpawnAgentHandlerV2::default() .handle(invocation( @@ -892,7 +898,7 @@ async fn multi_agent_v2_full_history_fork_accepts_explicit_service_tier() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let manager = thread_manager(); let root = manager .start_thread((*turn.config).clone()) @@ -960,6 +966,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { .expect("test config should allow feature update"); let turn = TurnContext { config: Arc::new(config), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::V2, ..turn }; @@ -1041,7 +1048,7 @@ async fn multi_agent_v2_spawn_requires_task_name() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let invocation = invocation( Arc::new(session), @@ -1075,7 +1082,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_items_field() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let invocation = invocation( Arc::new(session), @@ -1135,7 +1142,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -1232,7 +1239,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let err = SpawnAgentHandlerV2::default() .handle(invocation( @@ -1272,7 +1279,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let err = SpawnAgentHandlerV2::default() .handle(invocation( @@ -1312,7 +1319,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let err = SpawnAgentHandlerV2::default() .handle(invocation( @@ -1352,7 +1359,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let child_path = AgentPath::try_from("/root/worker").expect("agent path"); let child_thread_id = session @@ -1428,7 +1435,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let child_path = AgentPath::try_from("/root/worker").expect("agent path"); let child_thread_id = session @@ -1508,7 +1515,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa session.conversation_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -1602,7 +1609,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { session.conversation_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config.clone()); + set_turn_config(&mut turn, config.clone()); let researcher_path = AgentPath::from_string("/root/researcher".to_string()).expect("path"); let worker_path = AgentPath::from_string("/root/researcher/worker".to_string()).expect("path"); @@ -1689,7 +1696,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() { session.conversation_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -1753,7 +1760,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() { session.conversation_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -1809,7 +1816,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { session.conversation_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -1882,7 +1889,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() session.conversation_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -2017,7 +2024,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { session.conversation_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -2070,7 +2077,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { session.conversation_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -2150,7 +2157,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let output = SpawnAgentHandlerV2::default() .handle(invocation( @@ -2189,7 +2196,7 @@ async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let invocation = invocation( Arc::new(session), @@ -2399,7 +2406,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() { .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let parent_path = AgentPath::try_from("/root/parent").expect("agent path"); turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, @@ -2876,7 +2883,7 @@ async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -2961,7 +2968,7 @@ async fn multi_agent_v2_wait_agent_rejects_timeout_below_configured_min() { config.multi_agent_v2.min_wait_timeout_ms = 50; config.multi_agent_v2.max_wait_timeout_ms = 1_000; config.multi_agent_v2.default_wait_timeout_ms = 50; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let Err(err) = WaitAgentHandlerV2::default() .handle(invocation( @@ -2991,7 +2998,7 @@ async fn multi_agent_v2_wait_agent_accepts_explicit_timeout_at_configured_min() config.multi_agent_v2.min_wait_timeout_ms = 1; config.multi_agent_v2.max_wait_timeout_ms = 1_000; config.multi_agent_v2.default_wait_timeout_ms = 50; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let output = WaitAgentHandlerV2::default() .handle(invocation( @@ -3026,7 +3033,7 @@ async fn multi_agent_v2_wait_agent_uses_configured_default_timeout() { config.multi_agent_v2.min_wait_timeout_ms = 1; config.multi_agent_v2.max_wait_timeout_ms = 1_000; config.multi_agent_v2.default_wait_timeout_ms = 50; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3081,7 +3088,7 @@ async fn multi_agent_v2_wait_agent_allows_zero_configured_timeout() { config.multi_agent_v2.min_wait_timeout_ms = 0; config.multi_agent_v2.max_wait_timeout_ms = 0; config.multi_agent_v2.default_wait_timeout_ms = 0; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3121,7 +3128,7 @@ async fn multi_agent_v2_wait_agent_rejects_timeout_above_configured_max() { config.multi_agent_v2.min_wait_timeout_ms = 1; config.multi_agent_v2.max_wait_timeout_ms = 50; config.multi_agent_v2.default_wait_timeout_ms = 1; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let Err(err) = WaitAgentHandlerV2::default() .handle(invocation( @@ -3151,7 +3158,7 @@ async fn multi_agent_v2_wait_agent_accepts_explicit_timeout_at_configured_max() config.multi_agent_v2.min_wait_timeout_ms = 1; config.multi_agent_v2.max_wait_timeout_ms = 1; config.multi_agent_v2.default_wait_timeout_ms = 1; - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let output = WaitAgentHandlerV2::default() .handle(invocation( @@ -3359,7 +3366,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3453,7 +3460,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3534,7 +3541,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3625,7 +3632,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3714,7 +3721,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3762,7 +3769,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() { let (mut session, mut turn) = make_session_and_context().await; let mut config = (*turn.config).clone(); - config.agent_max_threads = Some(1); + config.multi_agent_v2.max_concurrent_threads_per_session = 2; config .features .enable(Feature::MultiAgentV2) @@ -3787,7 +3794,7 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() { .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - turn.config = Arc::new(config.clone()); + set_turn_config(&mut turn, config.clone()); let session = Arc::new(session); let turn = Arc::new(turn); @@ -3893,7 +3900,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); - turn.config = Arc::new(config); + set_turn_config(&mut turn, config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -4288,7 +4295,7 @@ async fn build_agent_resume_config_clears_base_instructions() { .set(AskForApproval::OnRequest) .expect("approval policy set"); - let config = build_agent_resume_config(&turn, /*child_depth*/ 0).expect("resume config"); + let config = build_agent_resume_config(&turn).expect("resume config"); let mut expected = (*turn.config).clone(); expected.base_instructions = None; 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 f767f88d78..a5309bdef0 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 @@ -108,7 +108,6 @@ async fn handle_spawn_agent( ) .await?; apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; - apply_spawn_agent_overrides(&mut config, child_depth); let spawn_source = thread_spawn_source( session.conversation_id, @@ -146,6 +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), }, ), ) diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 3c8abf8f16..d6e038ea78 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1,3 +1,5 @@ +use crate::agent::exceeds_thread_spawn_depth_limit; +use crate::agent::next_thread_spawn_depth; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::execute_spec::create_code_mode_tool; use crate::tools::context::ToolInvocation; @@ -61,6 +63,7 @@ use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_tools::DiscoverableTool; @@ -286,11 +289,18 @@ fn namespace_tools_enabled(turn_context: &TurnContext) -> bool { } fn multi_agent_v2_enabled(turn_context: &TurnContext) -> bool { - turn_context.features.get().enabled(Feature::MultiAgentV2) + turn_context.multi_agent_version == MultiAgentVersion::V2 } fn collab_tools_enabled(turn_context: &TurnContext) -> bool { - multi_agent_v2_enabled(turn_context) || turn_context.features.get().enabled(Feature::Collab) + match turn_context.multi_agent_version { + MultiAgentVersion::None => false, + MultiAgentVersion::V1 => !exceeds_thread_spawn_depth_limit( + next_thread_spawn_depth(&turn_context.session_source), + turn_context.config.agent_max_depth, + ), + MultiAgentVersion::V2 => true, + } } fn goal_tools_enabled(turn_context: &TurnContext) -> bool { @@ -303,6 +313,10 @@ fn goal_tools_enabled(turn_context: &TurnContext) -> bool { fn agent_jobs_tools_enabled(turn_context: &TurnContext) -> bool { turn_context.features.get().enabled(Feature::SpawnCsv) + && !exceeds_thread_spawn_depth_limit( + next_thread_spawn_depth(&turn_context.session_source), + turn_context.config.agent_max_depth, + ) } fn agent_jobs_worker_tools_enabled(turn_context: &TurnContext) -> bool { diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 3f832ed0ac..77a169ded8 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -15,6 +15,7 @@ use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ToolMode; use codex_protocol::openai_models::WebSearchToolType; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_tools::DiscoverablePluginInfo; @@ -215,6 +216,7 @@ fn set_feature(turn: &mut TurnContext, feature: Feature, enabled: bool) { .disable(feature) .expect("test feature should be disableable in config"); } + turn.multi_agent_version = config.multi_agent_version_from_features(); turn.config = Arc::new(config); turn.tool_mode = turn.model_info.tool_mode.unwrap_or_else(|| { if turn.config.features.enabled(Feature::CodeModeOnly) { @@ -528,6 +530,14 @@ async fn host_context_gates_goal_and_agent_job_tools() { normal_agent_job.assert_visible_contains(&["spawn_agents_on_csv"]); normal_agent_job.assert_visible_lacks(&["report_agent_job_result"]); + let agent_job_without_multi_agent = probe(|turn| { + set_feature(turn, Feature::SpawnCsv, /*enabled*/ true); + turn.multi_agent_version = MultiAgentVersion::None; + }) + .await; + agent_job_without_multi_agent.assert_visible_contains(&["spawn_agents_on_csv"]); + agent_job_without_multi_agent.assert_visible_lacks(&[MULTI_AGENT_V1_NAMESPACE, "spawn_agent"]); + let worker_agent_job = probe(|turn| { set_feature(turn, Feature::SpawnCsv, /*enabled*/ true); turn.session_source = @@ -864,6 +874,38 @@ async fn multi_agent_feature_selects_one_agent_tool_family() { ); } +#[tokio::test] +async fn multi_agent_tool_family_uses_locked_version_instead_of_features() { + let v1 = probe(|turn| { + set_feature(turn, Feature::MultiAgentV2, /*enabled*/ true); + turn.multi_agent_version = MultiAgentVersion::V1; + }) + .await; + v1.assert_visible_contains(&[MULTI_AGENT_V1_NAMESPACE]); + v1.assert_visible_lacks(&["spawn_agent", "send_message", "assign_task", "list_agents"]); + + let v2 = probe(|turn| { + set_feature(turn, Feature::Collab, /*enabled*/ false); + turn.multi_agent_version = MultiAgentVersion::V2; + }) + .await; + v2.assert_visible_contains(&["spawn_agent", "send_message", "assign_task", "list_agents"]); + v2.assert_visible_lacks(&[MULTI_AGENT_V1_NAMESPACE]); + + let disabled = probe(|turn| { + set_feature(turn, Feature::Collab, /*enabled*/ true); + turn.multi_agent_version = MultiAgentVersion::None; + }) + .await; + disabled.assert_visible_lacks(&[ + MULTI_AGENT_V1_NAMESPACE, + "spawn_agent", + "send_message", + "assign_task", + "list_agents", + ]); +} + #[tokio::test] async fn tool_mode_selector_overrides_feature_flags() { let direct = probe(|turn| { diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 84a2e18824..5302274153 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -75,6 +75,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; @@ -124,6 +125,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; diff --git a/codex-rs/core/tests/suite/rollout_list_find.rs b/codex-rs/core/tests/suite/rollout_list_find.rs index 4ff35449cb..cf11334105 100644 --- a/codex-rs/core/tests/suite/rollout_list_find.rs +++ b/codex-rs/core/tests/suite/rollout_list_find.rs @@ -189,6 +189,7 @@ async fn find_locates_rollout_file_written_by_recorder() -> std::io::Result<()> /*thread_source*/ None, BaseInstructions::default(), Vec::new(), + codex_protocol::protocol::MultiAgentVersion::None, ), ) .await?; diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index aecb9a392e..58267d5d5d 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -226,6 +226,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index a722f2e953..d7a040de01 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2453,6 +2453,16 @@ impl InitialHistory { } } + pub fn get_multi_agent_version(&self) -> Option { + match self { + InitialHistory::New | InitialHistory::Cleared => None, + InitialHistory::Resumed(resumed) => { + multi_agent_version_from_items(&resumed.history, Some(resumed.conversation_id)) + } + InitialHistory::Forked(items) => multi_agent_version_from_items(items, None), + } + } + pub fn get_resumed_session_sources(&self) -> Option<(SessionSource, Option)> { let meta = self.get_resumed_session_meta()?; Some((meta.source.clone(), meta.thread_source)) @@ -2723,6 +2733,29 @@ impl fmt::Display for InternalSessionSource { } } +fn multi_agent_version_from_items( + items: &[RolloutItem], + thread_id: Option, +) -> Option { + items.iter().rev().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) + if thread_id.is_none_or(|thread_id| meta_line.meta.id == thread_id) => + { + meta_line.meta.multi_agent_version + } + _ => None, + }) +} + +#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +pub enum MultiAgentVersion { + #[default] + None, + V1, + V2, +} + /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -2762,6 +2795,8 @@ pub struct SessionMeta { pub dynamic_tools: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub memory_mode: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multi_agent_version: Option, } impl Default for SessionMeta { @@ -2783,6 +2818,7 @@ impl Default for SessionMeta { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, } } } @@ -4008,6 +4044,29 @@ mod tests { ); } + #[test] + fn session_meta_multi_agent_version_distinguishes_missing_and_locked_values() { + let mut meta = SessionMeta::default(); + let serialized = serde_json::to_value(&meta).expect("serialize legacy session metadata"); + assert_eq!(serialized.get("multi_agent_version"), None); + let deserialized: SessionMeta = + serde_json::from_value(serialized).expect("deserialize legacy session metadata"); + assert_eq!(deserialized.multi_agent_version, None); + + for version in [ + MultiAgentVersion::None, + MultiAgentVersion::V1, + MultiAgentVersion::V2, + ] { + meta.multi_agent_version = Some(version); + let serialized = + serde_json::to_value(&meta).expect("serialize locked session metadata"); + let deserialized: SessionMeta = + serde_json::from_value(serialized).expect("deserialize locked session metadata"); + assert_eq!(deserialized.multi_agent_version, Some(version)); + } + } + #[test] fn inter_agent_communication_response_input_item_preserves_commentary_phase() { let communication = InterAgentCommunication { diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 7ae3843085..b53c958ae8 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -49,6 +49,7 @@ async fn extract_metadata_from_rollout_uses_session_meta() { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; let session_meta_line = SessionMetaLine { meta: session_meta, @@ -102,9 +103,11 @@ async fn extract_metadata_from_rollout_returns_latest_memory_mode() { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; let polluted_meta = SessionMeta { memory_mode: Some("polluted".to_string()), + multi_agent_version: None, ..session_meta.clone() }; let lines = vec![ @@ -363,6 +366,7 @@ fn write_rollout_in_sessions_with_cwd( base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; let session_meta_line = SessionMetaLine { meta: session_meta, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 94e3efe2ec..57010ff979 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -52,6 +52,7 @@ use codex_git_utils::collect_git_info; use codex_git_utils::get_git_repo_root; use codex_protocol::protocol::GitInfo as ProtocolGitInfo; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; @@ -87,6 +88,7 @@ pub enum RolloutRecorderParams { thread_source: Option, base_instructions: BaseInstructions, dynamic_tools: Vec, + multi_agent_version: MultiAgentVersion, }, Resume { path: PathBuf, @@ -163,6 +165,7 @@ impl RolloutRecorderParams { thread_source: Option, base_instructions: BaseInstructions, dynamic_tools: Vec, + multi_agent_version: MultiAgentVersion, ) -> Self { Self::Create { conversation_id, @@ -172,6 +175,7 @@ impl RolloutRecorderParams { thread_source, base_instructions, dynamic_tools, + multi_agent_version, } } @@ -661,6 +665,7 @@ impl RolloutRecorder { thread_source, base_instructions, dynamic_tools, + multi_agent_version, } => { let log_file_info = precompute_log_file_info(config, conversation_id)?; let path = log_file_info.path.clone(); @@ -696,6 +701,7 @@ impl RolloutRecorder { Some(dynamic_tools) }, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), + multi_agent_version: Some(multi_agent_version), }; (None, Some(log_file_info), path, Some(session_meta)) diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index a9ee6ca9aa..8771c5f5f7 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -99,6 +99,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; @@ -375,6 +376,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< /*thread_source*/ None, BaseInstructions::default(), Vec::new(), + codex_protocol::protocol::MultiAgentVersion::None, ), ) .await?; @@ -456,6 +458,7 @@ async fn persist_reports_filesystem_error_and_retries_buffered_items() -> std::i /*thread_source*/ None, BaseInstructions::default(), Vec::new(), + codex_protocol::protocol::MultiAgentVersion::None, ), ) .await?; diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index 91b475caff..97496d394e 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -41,6 +41,7 @@ fn write_rollout_with_metadata(path: &Path, thread_id: ThreadId) -> std::io::Res base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }), diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 73431b0412..e489eaba59 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -1274,6 +1274,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }), diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index f64e6f7fed..0d15d46a52 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -336,6 +336,7 @@ mod tests { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }), @@ -494,6 +495,7 @@ mod tests { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }), diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 9522b6176d..4cc7ee4dc1 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1332,6 +1332,7 @@ mod tests { base_instructions: None, dynamic_tools: None, memory_mode: Some("polluted".to_string()), + multi_agent_version: None, }, git: None, })]; @@ -1392,6 +1393,7 @@ mod tests { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: Some(GitInfo { commit_hash: Some(codex_git_utils::GitSha::new("rollout-sha")), diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 02e64a1ff9..008fee6cbe 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -10,7 +10,11 @@ use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadMemoryMode; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; @@ -20,6 +24,7 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::SetMultiAgentVersionIfUnsetParams; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; @@ -43,6 +48,105 @@ mod tests { use crate::ListTurnsParams; use crate::SortDirection; use crate::StoredTurnItemsView; + use crate::ThreadEventPersistenceMode; + use crate::ThreadPersistenceMetadata; + use codex_protocol::models::BaseInstructions; + use codex_protocol::protocol::SessionSource; + use pretty_assertions::assert_eq; + + #[tokio::test] + async fn create_thread_records_multi_agent_version() { + let store = InMemoryThreadStore::default(); + let thread_id = ThreadId::new(); + store + .create_thread(CreateThreadParams { + thread_id, + forked_from_id: None, + source: SessionSource::Exec, + thread_source: None, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + multi_agent_version: MultiAgentVersion::V2, + metadata: ThreadPersistenceMetadata { + cwd: None, + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + event_persistence_mode: ThreadEventPersistenceMode::Limited, + }) + .await + .expect("create thread should succeed"); + + let history = store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: true, + }) + .await + .expect("created history should load"); + let versions = history + .items + .into_iter() + .filter_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.multi_agent_version, + _ => None, + }) + .collect::>(); + assert_eq!(versions, vec![MultiAgentVersion::V2]); + } + + #[tokio::test] + async fn concurrent_multi_agent_version_seeds_return_the_first_winner() { + let store = InMemoryThreadStore::default(); + let thread_id = ThreadId::new(); + store.state.lock().await.histories.insert( + thread_id, + vec![RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + ..SessionMeta::default() + }, + git: None, + })], + ); + + let (left, right) = tokio::join!( + store.set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V1, + include_archived: true, + }), + store.set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V2, + include_archived: true, + }), + ); + let left = left.expect("first seed should succeed"); + let right = right.expect("second seed should return the stored winner"); + assert_eq!(left, right); + assert!(matches!( + left, + MultiAgentVersion::V1 | MultiAgentVersion::V2 + )); + + let history = store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: true, + }) + .await + .expect("seeded history should load"); + let versions = history + .items + .into_iter() + .filter_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.multi_agent_version, + _ => None, + }) + .collect::>(); + assert_eq!(versions, vec![left]); + } #[tokio::test] async fn default_turn_pagination_methods_return_unsupported() { @@ -99,6 +203,7 @@ fn stores_guard() -> MutexGuard<'static, HashMap ThreadStoreResult<()> { let mut state = self.state.lock().await; state.calls.create_thread += 1; - state.histories.entry(params.thread_id).or_default(); + let session_meta = SessionMeta { + id: params.thread_id, + forked_from_id: params.forked_from_id, + cwd: params.metadata.cwd.clone().unwrap_or_default(), + agent_nickname: params.source.get_nickname(), + agent_role: params.source.get_agent_role(), + agent_path: params.source.get_agent_path().map(Into::into), + source: params.source.clone(), + thread_source: params.thread_source, + model_provider: Some(params.metadata.model_provider.clone()), + base_instructions: Some(params.base_instructions.clone()), + dynamic_tools: (!params.dynamic_tools.is_empty()).then(|| params.dynamic_tools.clone()), + memory_mode: matches!(params.metadata.memory_mode, ThreadMemoryMode::Disabled) + .then_some("disabled".to_string()), + multi_agent_version: Some(params.multi_agent_version), + ..SessionMeta::default() + }; + state + .histories + .entry(params.thread_id) + .or_default() + .push(RolloutItem::SessionMeta(SessionMetaLine { + meta: session_meta, + git: None, + })); state.created_threads.insert(params.thread_id, params); Ok(()) } @@ -179,6 +308,44 @@ impl ThreadStore for InMemoryThreadStore { Ok(()) } + async fn set_multi_agent_version_if_unset( + &self, + params: SetMultiAgentVersionIfUnsetParams, + ) -> ThreadStoreResult { + let mut state = self.state.lock().await; + state.calls.set_multi_agent_version_if_unset += 1; + let history = + state + .histories + .get_mut(¶ms.thread_id) + .ok_or(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + })?; + if let Some(multi_agent_version) = history.iter().rev().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == params.thread_id => { + meta_line.meta.multi_agent_version + } + _ => None, + }) { + return Ok(multi_agent_version); + } + let mut session_meta = history + .iter() + .find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == params.thread_id => { + Some(meta_line.clone()) + } + _ => None, + }) + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("thread {} does not have session metadata", params.thread_id), + })?; + session_meta.git = None; + session_meta.meta.multi_agent_version = Some(params.multi_agent_version); + history.push(RolloutItem::SessionMeta(session_meta)); + Ok(params.multi_agent_version) + } + async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> { let mut state = self.state.lock().await; state.calls.append_items += 1; diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index aec8270a53..53cde2715d 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -35,6 +35,7 @@ pub use types::ReadThreadByRolloutPathParams; pub use types::ReadThreadParams; pub use types::ResumeThreadParams; pub use types::SearchThreadsParams; +pub use types::SetMultiAgentVersionIfUnsetParams; pub use types::SortDirection; pub use types::StoredThread; pub use types::StoredThreadHistory; diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index a71bb48b0f..a2735261c8 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -35,6 +35,7 @@ pub(super) async fn create_thread( params.thread_source, params.base_instructions, params.dynamic_tools, + params.multi_agent_version, ), ) .await diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index e26c314e73..aa2a9819d6 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -5,6 +5,7 @@ mod list_threads; mod live_writer; mod read_thread; mod search_threads; +mod set_multi_agent_version; mod unarchive_thread; mod update_thread_metadata; @@ -30,6 +31,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::SearchThreadsParams; +use crate::SetMultiAgentVersionIfUnsetParams; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadPage; @@ -38,6 +40,7 @@ use crate::ThreadStore; use crate::ThreadStoreError; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; +use codex_protocol::protocol::MultiAgentVersion; /// Local filesystem/SQLite-backed implementation of [`ThreadStore`]. /// @@ -56,6 +59,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, + multi_agent_version_seed_lock: Arc>, state_db: Option, } @@ -95,6 +99,7 @@ impl LocalThreadStore { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), + multi_agent_version_seed_lock: Arc::new(Mutex::new(())), state_db, } } @@ -180,6 +185,13 @@ impl ThreadStore for LocalThreadStore { live_writer::resume_thread(self, params).await } + async fn set_multi_agent_version_if_unset( + &self, + params: SetMultiAgentVersionIfUnsetParams, + ) -> ThreadStoreResult { + set_multi_agent_version::set_multi_agent_version_if_unset(self, params).await + } + async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> { live_writer::append_items(self, params).await } @@ -301,6 +313,7 @@ mod tests { use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::UserMessageEvent; + use pretty_assertions::assert_eq; use tempfile::TempDir; use super::*; @@ -817,6 +830,49 @@ mod tests { assert!(err.to_string().contains("requires a cwd")); } + #[tokio::test] + async fn set_multi_agent_version_if_unset_appends_one_lock_to_legacy_rollout() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let uuid = uuid::Uuid::from_u128(408); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let rollout_path = + write_session_file(home.path(), "2025-01-04T12-00-00", uuid).expect("session file"); + + let first = store + .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V1, + include_archived: true, + }) + .await + .expect("first seed should succeed"); + let second = store + .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V2, + include_archived: true, + }) + .await + .expect("second seed should return the stored winner"); + assert_eq!( + (first, second), + (MultiAgentVersion::V1, MultiAgentVersion::V1) + ); + + let (items, _, _) = RolloutRecorder::load_rollout_items(&rollout_path) + .await + .expect("load rollout items"); + let versions = items + .into_iter() + .filter_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.multi_agent_version, + _ => None, + }) + .collect::>(); + assert_eq!(versions, vec![MultiAgentVersion::V1]); + } + #[tokio::test] async fn load_history_uses_live_writer_rollout_path() { let home = TempDir::new().expect("temp dir"); @@ -1028,6 +1084,7 @@ mod tests { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: codex_protocol::protocol::MultiAgentVersion::None, metadata: thread_metadata(), event_persistence_mode: ThreadEventPersistenceMode::Limited, } diff --git a/codex-rs/thread-store/src/local/set_multi_agent_version.rs b/codex-rs/thread-store/src/local/set_multi_agent_version.rs new file mode 100644 index 0000000000..ee2711c569 --- /dev/null +++ b/codex-rs/thread-store/src/local/set_multi_agent_version.rs @@ -0,0 +1,89 @@ +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::RolloutItem; +use codex_rollout::append_rollout_item_to_path; +use codex_rollout::read_session_meta_line; + +use super::LocalThreadStore; +use super::live_writer; +use super::read_thread; +use crate::AppendThreadItemsParams; +use crate::ReadThreadParams; +use crate::SetMultiAgentVersionIfUnsetParams; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +pub(super) async fn set_multi_agent_version_if_unset( + store: &LocalThreadStore, + params: SetMultiAgentVersionIfUnsetParams, +) -> ThreadStoreResult { + let _guard = store.multi_agent_version_seed_lock.lock().await; + let has_live_writer = live_writer::rollout_path(store, params.thread_id) + .await + .is_ok(); + if has_live_writer { + live_writer::persist_thread(store, params.thread_id).await?; + live_writer::flush_thread(store, params.thread_id).await?; + } + + let thread = read_thread::read_thread( + store, + ReadThreadParams { + thread_id: params.thread_id, + include_archived: params.include_archived, + include_history: true, + }, + ) + .await?; + let rollout_path = thread + .rollout_path + .ok_or_else(|| ThreadStoreError::Internal { + message: format!("thread {} does not have a rollout path", params.thread_id), + })?; + let history = thread.history.ok_or_else(|| ThreadStoreError::Internal { + message: format!("failed to load history for thread {}", params.thread_id), + })?; + if let Some(multi_agent_version) = history.items.iter().rev().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == params.thread_id => { + meta_line.meta.multi_agent_version + } + _ => None, + }) { + return Ok(multi_agent_version); + } + + let mut session_meta = + read_session_meta_line(&rollout_path) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to set thread multi-agent version: {err}"), + })?; + if session_meta.meta.id != params.thread_id { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to set thread multi-agent version: rollout session metadata id mismatch: expected {}, found {}", + params.thread_id, session_meta.meta.id + ), + }); + } + session_meta.git = None; + session_meta.meta.multi_agent_version = Some(params.multi_agent_version); + let item = RolloutItem::SessionMeta(session_meta); + if has_live_writer { + live_writer::append_items( + store, + AppendThreadItemsParams { + thread_id: params.thread_id, + items: vec![item], + }, + ) + .await?; + } else { + append_rollout_item_to_path(&rollout_path, &item) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to set thread multi-agent version: {err}"), + })?; + } + + Ok(params.multi_agent_version) +} diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 7276a980da..66f9a38774 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -14,6 +14,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::SearchThreadsParams; +use crate::SetMultiAgentVersionIfUnsetParams; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadPage; @@ -22,6 +23,7 @@ use crate::ThreadStoreError; use crate::ThreadStoreResult; use crate::TurnPage; use crate::UpdateThreadMetadataParams; +use codex_protocol::protocol::MultiAgentVersion; /// Storage-neutral thread persistence boundary. #[async_trait] @@ -35,6 +37,12 @@ pub trait ThreadStore: Any + Send + Sync { /// Reopens an existing thread for live appends. async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()>; + /// Persists a multi-agent runtime for a legacy thread if one has not been stored yet. + async fn set_multi_agent_version_if_unset( + &self, + params: SetMultiAgentVersionIfUnsetParams, + ) -> ThreadStoreResult; + /// Appends canonical rollout items to a live thread. /// /// This is the raw history API. It does not infer metadata from item contents. Callers that diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 723f08c116..8b51cb155b 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -9,6 +9,7 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadMemoryMode as MemoryMode; @@ -65,6 +66,17 @@ pub struct ThreadPersistenceMetadata { pub memory_mode: MemoryMode, } +/// Parameters for locking a legacy thread to a multi-agent runtime. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SetMultiAgentVersionIfUnsetParams { + /// Thread id whose metadata should be initialized. + pub thread_id: ThreadId, + /// Version to persist if the thread does not already have one. + pub multi_agent_version: MultiAgentVersion, + /// Whether archived threads are eligible. + pub include_archived: bool, +} + /// Parameters required to create a persisted thread. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct CreateThreadParams { @@ -82,6 +94,8 @@ pub struct CreateThreadParams { pub base_instructions: BaseInstructions, /// Dynamic tools available to the thread at startup. pub dynamic_tools: Vec, + /// Multi-agent runtime selected when the thread was created. + pub multi_agent_version: MultiAgentVersion, /// Metadata captured for the newly created thread. pub metadata: ThreadPersistenceMetadata, /// Whether persistence should include the extended event surface.