diff --git a/codex-rs/app-server/tests/common/models_cache.rs b/codex-rs/app-server/tests/common/models_cache.rs index 56a9378662..666f6ed7cd 100644 --- a/codex-rs/app-server/tests/common/models_cache.rs +++ b/codex-rs/app-server/tests/common/models_cache.rs @@ -54,6 +54,7 @@ fn preset_to_info(preset: &ModelPreset, priority: i32) -> ModelInfo { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, } } 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..fc32716a9b 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: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_fork.rs b/codex-rs/app-server/tests/suite/v2/thread_fork.rs index 5c4c637f4b..7fa3942979 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -29,6 +29,10 @@ use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; use codex_config::types::AuthCredentialsStoreMode; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; +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 pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -81,6 +85,9 @@ async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { "expected original rollout to exist at {}", original_path.display() ); + let mut session_meta = read_session_meta_line(&original_path).await?; + session_meta.meta.multi_agent_version = Some(MultiAgentVersion::V1); + append_rollout_item_to_path(&original_path, &RolloutItem::SessionMeta(session_meta)).await?; let original_contents = std::fs::read_to_string(&original_path)?; let mut mcp = McpProcess::new(codex_home.path()).await?; 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..ce375fc330 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: 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..65e0571326 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -63,6 +63,8 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ImageGenerationEndEvent; use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::McpToolCallEndEvent; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource as RolloutSessionSource; @@ -74,6 +76,8 @@ use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; +use codex_rollout::append_rollout_item_to_path; +use codex_rollout::read_session_meta_line; use codex_state::StateRuntime; use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::responses; @@ -1805,6 +1809,7 @@ stream_max_retries = 0 base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }; std::fs::write( &rollout_path, @@ -2013,7 +2018,7 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - let rollout = setup_rollout_fixture(codex_home.path(), &server.uri())?; + let rollout = setup_rollout_fixture(codex_home.path(), &server.uri()).await?; let thread_id = rollout.conversation_id.clone(); let mut mcp = McpProcess::new(codex_home.path()).await?; @@ -3051,7 +3056,7 @@ async fn thread_resume_with_overrides_defers_updated_at_until_turn_start() -> Re async fn thread_resume_fails_when_required_mcp_server_fails_to_initialize() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - let rollout = setup_rollout_fixture(codex_home.path(), &server.uri())?; + let rollout = setup_rollout_fixture(codex_home.path(), &server.uri()).await?; create_config_toml_with_required_broken_mcp(codex_home.path(), &server.uri())?; let mut mcp = McpProcess::new(codex_home.path()).await?; @@ -3614,7 +3619,7 @@ struct RolloutFixture { before_modified: std::time::SystemTime, } -fn setup_rollout_fixture(codex_home: &Path, server_uri: &str) -> Result { +async fn setup_rollout_fixture(codex_home: &Path, server_uri: &str) -> Result { create_config_toml(codex_home, server_uri)?; let preview = "Saved user message"; @@ -3631,6 +3636,10 @@ fn setup_rollout_fixture(codex_home: &Path, server_uri: &str) -> Result Result<()> { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/codex-api/tests/models_integration.rs b/codex-rs/codex-api/tests/models_integration.rs index 742b39d5fd..9ff273211f 100644 --- a/codex-rs/codex-api/tests/models_integration.rs +++ b/codex-rs/codex-api/tests/models_integration.rs @@ -100,6 +100,7 @@ async fn models_client_hits_models_endpoint() { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, }], }; diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1f58fb437e..d9a84455f5 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -5,14 +5,15 @@ use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; +use crate::config::Config; use crate::session::emit_subagent_session_started; +use crate::session::resolve_multi_agent_version; use crate::session_prefix::format_subagent_context_line; use crate::session_prefix::format_subagent_notification_message; 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; @@ -23,6 +24,7 @@ use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; 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; @@ -58,6 +60,12 @@ pub(crate) struct SpawnAgentOptions { pub(crate) environments: Option>, } +struct SpawnAgentThreadInheritance { + shell_snapshot: Option>, + exec_policy: Option>, + multi_agent_version: MultiAgentVersion, +} + #[derive(Clone, Debug)] pub(crate) struct LiveAgent { pub(crate) thread_id: ThreadId, @@ -80,10 +88,7 @@ fn default_agent_nickname_list() -> Vec<&'static str> { .collect() } -fn agent_nickname_candidates( - config: &crate::config::Config, - role_name: Option<&str>, -) -> Vec { +fn agent_nickname_candidates(config: &Config, role_name: Option<&str>) -> Vec { let role_name = role_name.unwrap_or(DEFAULT_ROLE_NAME); if let Some(candidates) = resolve_role_config(config, role_name).and_then(|role| role.nickname_candidates.clone()) @@ -185,14 +190,32 @@ impl AgentControl { #[cfg(test)] pub(crate) async fn spawn_agent( &self, - config: crate::config::Config, + config: Config, initial_operation: Op, session_source: Option, ) -> CodexResult { + let state = self.upgrade()?; + let multi_agent_version = + if let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + .. + })) = session_source.as_ref() + { + state + .get_thread(*parent_thread_id) + .await? + .multi_agent_version() + .ok_or_else(|| { + CodexErr::InvalidRequest("multi-agent version is unresolved".to_string()) + })? + } else { + config.multi_agent_version_from_features() + }; let spawned_agent = Box::pin(self.spawn_agent_internal( config, initial_operation, session_source, + multi_agent_version, SpawnAgentOptions::default(), )) .await?; @@ -202,30 +225,44 @@ impl AgentControl { /// Spawn an agent thread with some metadata. pub(crate) async fn spawn_agent_with_metadata( &self, - config: crate::config::Config, + config: Config, initial_operation: Op, session_source: Option, + multi_agent_version: MultiAgentVersion, options: SpawnAgentOptions, // TODO(jif) drop with new fork. ) -> CodexResult { - Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options)) - .await + Box::pin(self.spawn_agent_internal( + config, + initial_operation, + session_source, + multi_agent_version, + options, + )) + .await } async fn spawn_agent_internal( &self, - config: crate::config::Config, + config: Config, initial_operation: Op, session_source: Option, + multi_agent_version: MultiAgentVersion, options: SpawnAgentOptions, ) -> CodexResult { let state = self.upgrade()?; - 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()) - .await; - let inherited_exec_policy = self - .inherited_exec_policy_for_source(&state, session_source.as_ref(), &config) - .await; + let agent_max_threads = config + .effective_agent_max_threads(multi_agent_version) + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; + let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; + let inheritance = SpawnAgentThreadInheritance { + shell_snapshot: self + .inherited_shell_snapshot_for_source(&state, session_source.as_ref()) + .await, + exec_policy: self + .inherited_exec_policy_for_source(&state, session_source.as_ref(), &config) + .await, + multi_agent_version, + }; let (session_source, mut agent_metadata) = match session_source { Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, @@ -250,19 +287,18 @@ impl AgentControl { let notification_source = session_source.clone(); // The same `AgentControl` is sent to spawn the thread. - let new_thread = match (session_source, options.fork_mode.as_ref()) { - (Some(session_source), Some(_)) => { + let new_thread = match (session_source, options.fork_mode.as_ref(), inheritance) { + (Some(session_source), Some(_), inheritance) => { Box::pin(self.spawn_forked_thread( &state, config, session_source, &options, - inherited_shell_snapshot, - inherited_exec_policy, + inheritance, )) .await? } - (Some(session_source), None) => { + (Some(session_source), None, inheritance) => { Box::pin(state.spawn_new_thread_with_source( config.clone(), self.clone(), @@ -272,13 +308,14 @@ impl AgentControl { /*thread_source*/ Some(ThreadSource::Subagent), /*persist_extended_history*/ false, /*metrics_service_name*/ None, - inherited_shell_snapshot, - inherited_exec_policy, + inheritance.shell_snapshot, + inheritance.exec_policy, options.environments.clone(), + Some(inheritance.multi_agent_version), )) .await? } - (None, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, + (None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, }; agent_metadata.agent_id = Some(new_thread.thread_id); reservation.commit(agent_metadata.clone()); @@ -341,7 +378,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() != Some(MultiAgentVersion::V2) { let child_reference = agent_metadata .agent_path .as_ref() @@ -365,12 +402,16 @@ impl AgentControl { async fn spawn_forked_thread( &self, state: &Arc, - config: crate::config::Config, + config: Config, session_source: SessionSource, options: &SpawnAgentOptions, - inherited_shell_snapshot: Option>, - inherited_exec_policy: Option>, + inheritance: SpawnAgentThreadInheritance, ) -> CodexResult { + let SpawnAgentThreadInheritance { + shell_snapshot: inherited_shell_snapshot, + exec_policy: inherited_exec_policy, + multi_agent_version, + } = inheritance; if options.fork_parent_spawn_call_id.is_none() { return Err(CodexErr::Fatal( "spawn_agent fork requires a parent spawn call id".to_string(), @@ -420,7 +461,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 @@ -438,7 +479,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(), @@ -474,7 +515,7 @@ impl AgentControl { } } if preserve_reference_context_item - && config.features.enabled(Feature::MultiAgentV2) + && multi_agent_version == MultiAgentVersion::V2 && config.multi_agent_v2.usage_hint_enabled && let Some(subagent_usage_hint_text) = config.multi_agent_v2.subagent_usage_hint_text.clone() @@ -499,6 +540,7 @@ impl AgentControl { inherited_shell_snapshot, inherited_exec_policy, options.environments.clone(), + Some(multi_agent_version), ) .await } @@ -506,7 +548,7 @@ impl AgentControl { /// Resume an existing agent thread from a recorded rollout file. pub(crate) async fn resume_agent_from_rollout( &self, - config: crate::config::Config, + config: Config, thread_id: ThreadId, session_source: SessionSource, ) -> CodexResult { @@ -581,20 +623,48 @@ impl AgentControl { async fn resume_single_agent_from_rollout( &self, - mut config: crate::config::Config, + config: Config, 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 mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; + let inherited_multi_agent_version = + if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) = &session_source + { + match state.get_thread(*parent_thread_id).await { + Ok(parent_thread) => parent_thread.multi_agent_version(), + Err(_) => None, + } + } else { + None + }; + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: true, + }) + .await?; + let history = stored_thread + .history + .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))? + .items; + let initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history, + rollout_path: stored_thread.rollout_path, + }); + let parent_thread_id = stored_thread.parent_thread_id; + let multi_agent_version = + resolve_multi_agent_version(&initial_history, inherited_multi_agent_version) + .unwrap_or_else(|| config.multi_agent_version_from_features()); + let agent_max_threads = config + .effective_agent_max_threads(multi_agent_version) + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; + let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, @@ -631,32 +701,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; @@ -665,7 +720,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() != Some(MultiAgentVersion::V2) { let child_reference = agent_metadata .agent_path .as_ref() @@ -1050,12 +1105,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() == Some(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; }; @@ -1091,7 +1147,7 @@ impl AgentControl { fn prepare_thread_spawn( &self, reservation: &mut crate::agent::registry::SpawnReservation, - config: &crate::config::Config, + config: &Config, parent_thread_id: ThreadId, depth: i32, agent_path: Option, @@ -1153,7 +1209,7 @@ impl AgentControl { &self, state: &Arc, session_source: Option<&SessionSource>, - child_config: &crate::config::Config, + child_config: &Config, ) -> Option> { let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, .. diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3e0242c037..4b772d74fe 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -668,6 +668,9 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::FullHistory), @@ -737,6 +740,9 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::FullHistory), @@ -879,6 +885,9 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id), fork_mode: Some(SpawnAgentForkMode::FullHistory), @@ -949,6 +958,9 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::FullHistory), @@ -1066,6 +1078,9 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), @@ -1177,6 +1192,9 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id), fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), @@ -1287,6 +1305,9 @@ async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { agent_nickname: None, agent_role: None, })), + parent_thread + .multi_agent_version() + .expect("parent thread should have a multi-agent version"), SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id), fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), @@ -1587,9 +1608,15 @@ async fn spawn_child_completion_notifies_parent_history() { #[tokio::test] async fn multi_agent_v2_completion_ignores_dead_direct_parent() { let harness = AgentControlHarness::new().await; - let (root_thread_id, root_thread) = harness.start_thread().await; let mut config = harness.config.clone(); let _ = config.features.enable(Feature::MultiAgentV2); + let root = harness + .manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + let root_thread_id = root.thread_id; + let root_thread = root.thread; let worker_path = AgentPath::root().join("worker_a").expect("worker path"); let worker_thread_id = harness .control diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 44a8c0ac6b..ff08a61355 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; @@ -104,6 +105,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(), + inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) .or_cancel(&cancel_token) .await??; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 7a6bc69c5c..cd90829c0c 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 fn multi_agent_version(&self) -> Option { + self.codex.session.multi_agent_version() + } + /// 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..05cbb9828f 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -86,6 +86,7 @@ use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::NetworkAccess; use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::SandboxPolicy; @@ -9801,7 +9802,13 @@ non_code_mode_only = true assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 2500); assert_eq!(config.multi_agent_v2.max_wait_timeout_ms, 120000); assert_eq!(config.multi_agent_v2.default_wait_timeout_ms, 30000); - assert_eq!(config.agent_max_threads, Some(4)); + assert_eq!( + ( + config.agent_max_threads, + config.effective_agent_max_threads(MultiAgentVersion::V2)? + ), + (None, Some(4)) + ); assert!(!config.multi_agent_v2.usage_hint_enabled); assert_eq!( config.multi_agent_v2.usage_hint_text.as_deref(), @@ -9845,7 +9852,13 @@ enabled = true assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 10_000); assert_eq!(config.multi_agent_v2.max_wait_timeout_ms, 3_600_000); assert_eq!(config.multi_agent_v2.default_wait_timeout_ms, 30_000); - assert_eq!(config.agent_max_threads, Some(3)); + assert_eq!( + ( + config.agent_max_threads, + config.effective_agent_max_threads(MultiAgentVersion::V2)? + ), + (None, Some(3)) + ); assert_eq!( config.multi_agent_v2.root_agent_usage_hint_text.as_deref(), Some(DEFAULT_MULTI_AGENT_V2_ROOT_AGENT_USAGE_HINT_TEXT) @@ -9912,17 +9925,19 @@ max_threads = 3 "#, )?; - let err = ConfigBuilder::without_managed_config_for_tests() + let 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 + .effective_agent_max_threads(MultiAgentVersion::V2) .expect_err("agents.max_threads should conflict with multi_agent_v2"); assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); assert_eq!( err.to_string(), - "agents.max_threads cannot be set when multi_agent_v2 is enabled" + "agents.max_threads cannot be set when the multi-agent runtime is v2" ); Ok(()) @@ -10181,7 +10196,13 @@ max_concurrent_threads_per_session = 1 .await?; assert_eq!(config.multi_agent_v2.max_concurrent_threads_per_session, 1); - assert_eq!(config.agent_max_threads, Some(0)); + assert_eq!( + ( + config.agent_max_threads, + config.effective_agent_max_threads(MultiAgentVersion::V2)? + ), + (None, Some(0)) + ); Ok(()) } diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ff666b51e9..051064f279 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; @@ -823,7 +824,7 @@ pub struct Config { /// Token budget applied when storing tool/function outputs in the context manager. pub tool_output_token_limit: Option, - /// Maximum number of agent threads that can be open concurrently. + /// User-configured maximum number of agent threads that can be open concurrently. pub 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 +1282,40 @@ 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::Disabled + } + } + + pub(crate) fn effective_agent_max_threads( + &self, + multi_agent_version: MultiAgentVersion, + ) -> std::io::Result> { + match multi_agent_version { + MultiAgentVersion::V2 => { + if self.agent_max_threads.is_some() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "agents.max_threads cannot be set when the multi-agent runtime is v2", + )); + } + Ok(Some( + self.multi_agent_v2 + .max_concurrent_threads_per_session + .saturating_sub(1), + )) + } + MultiAgentVersion::Disabled | MultiAgentVersion::V1 => { + Ok(self.agent_max_threads.or(DEFAULT_AGENT_MAX_THREADS)) + } + } + } + pub fn legacy_sandbox_policy(&self) -> SandboxPolicy { self.permissions.legacy_sandbox_policy(self.cwd.as_path()) } @@ -3063,29 +3098,13 @@ 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 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 - }; + let agent_max_threads = cfg.agents.as_ref().and_then(|agents| agents.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", + )); + } let agent_max_depth = cfg .agents .as_ref() 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/mod.rs b/codex-rs/core/src/session/mod.rs index 12254f192d..5babd5e64b 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; @@ -137,6 +138,7 @@ use codex_thread_store::LiveThreadInitGuard; use codex_thread_store::LocalThreadStore; use codex_thread_store::ReadThreadParams; use codex_thread_store::ResumeThreadParams; +use codex_thread_store::SetMultiAgentVersionIfUnsetParams; use codex_thread_store::ThreadEventPersistenceMode; use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; @@ -420,6 +422,21 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) inherited_multi_agent_version: Option, +} + +pub(crate) fn resolve_multi_agent_version( + conversation_history: &InitialHistory, + inherited_multi_agent_version: Option, +) -> Option { + conversation_history + .get_multi_agent_version() + .or(inherited_multi_agent_version) + .or_else(|| match conversation_history { + InitialHistory::New | InitialHistory::Cleared => None, + // Threads created before runtime metadata existed keep the legacy V1 tool surface. + InitialHistory::Resumed(_) | InitialHistory::Forked(_) => Some(MultiAgentVersion::V1), + }) } pub(crate) const INITIAL_SUBMIT_ID: &str = ""; @@ -482,18 +499,11 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + inherited_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); - } - let primary_environment = environment_selections.primary_environment(); let mut user_instruction_warnings = Vec::new(); let user_instructions = AgentsMdManager::new(&config) @@ -544,6 +554,13 @@ impl Codex { let model_info = models_manager .get_model_info(model.as_str(), &config.to_models_manager_config()) .await; + let multi_agent_version = + resolve_multi_agent_version(&conversation_history, inherited_multi_agent_version); + let _ = config + .effective_agent_max_threads( + multi_agent_version.unwrap_or_else(|| config.multi_agent_version_from_features()), + ) + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; let base_instructions = config .base_instructions .clone() @@ -629,6 +646,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_provider, + multi_agent_version, ) .await .map_err(|e| { @@ -1652,7 +1670,7 @@ impl Session { turn_context: &TurnContext, msg: &EventMsg, ) { - if !self.enabled(Feature::MultiAgentV2) { + if turn_context.multi_agent_version != Some(MultiAgentVersion::V2) { return; } @@ -2623,6 +2641,56 @@ impl Session { state.session_configuration.collaboration_mode.clone() } + pub(crate) fn multi_agent_version(&self) -> Option { + self.multi_agent_version.get().copied() + } + + pub(crate) fn set_multi_agent_version_if_unset( + &self, + multi_agent_version: MultiAgentVersion, + ) -> MultiAgentVersion { + *self.multi_agent_version.get_or_init(|| multi_agent_version) + } + + pub(crate) async fn resolve_multi_agent_version_for_model( + &self, + model_info: &ModelInfo, + config: &Config, + ) -> MultiAgentVersion { + if let Some(multi_agent_version) = self.multi_agent_version() { + return multi_agent_version; + } + + let selected = model_info + .multi_agent_version + .unwrap_or_else(|| config.multi_agent_version_from_features()); + + let stored = if config.ephemeral { + selected + } else { + match self + .services + .thread_store + .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id: self.thread_id(), + multi_agent_version: selected, + }) + .await + { + Ok(multi_agent_version) => multi_agent_version, + Err(err) => { + warn!( + "failed to set multi-agent version for thread {}: {err}", + self.thread_id() + ); + selected + } + } + }; + + self.set_multi_agent_version_if_unset(stored) + } + 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..e4423b29ed 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 != Some(MultiAgentVersion::V2) { return None; } diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index c476436560..b8d54a1aa1 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: Some(MultiAgentVersion::Disabled), 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..937b5a08d7 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -9,8 +9,10 @@ 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 std::sync::OnceLock; use tokio::sync::Semaphore; /// Context for an initialized model agent @@ -29,6 +31,7 @@ pub(crate) struct Session { /// The set of enabled features should be invariant for the lifetime of the /// session. pub(super) features: ManagedFeatures, + pub(super) multi_agent_version: OnceLock, pub(super) pending_mcp_server_refresh_config: Mutex>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, @@ -504,6 +507,7 @@ impl Session { thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, + multi_agent_version: Option, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -518,6 +522,8 @@ impl Session { .parent_thread_id .or_else(|| initial_history.get_resumed_parent_thread_id()); session_configuration.parent_thread_id = parent_thread_id; + let multi_agent_version = multi_agent_version.map(OnceLock::from).unwrap_or_default(); + let initial_multi_agent_version = multi_agent_version.get().copied(); let event_persistence_mode = if session_configuration.persist_extended_history { ThreadEventPersistenceMode::Extended @@ -564,6 +570,7 @@ impl Session { text: session_configuration.base_instructions.clone(), }, dynamic_tools: session_configuration.dynamic_tools.clone(), + multi_agent_version: initial_multi_agent_version, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -1070,6 +1077,7 @@ impl Session { state: Mutex::new(state), managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), + multi_agent_version, pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 953f830f4c..2d5db9f23e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -112,6 +112,7 @@ use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::NetworkApprovalProtocol; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -121,6 +122,8 @@ use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::RealtimeVoicesList; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SkillScope; use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadGoalStatus; @@ -176,6 +179,7 @@ use serde::Deserialize; use serde_json::json; use std::path::PathBuf; use std::sync::Arc; +use std::sync::OnceLock; use std::time::Duration as StdDuration; mod guardian_tests; @@ -1645,6 +1649,56 @@ async fn record_initial_history_reconstructs_resumed_transcript() { assert_eq!(expected, history.raw_items()); } +#[test] +fn resolve_multi_agent_version_handles_unset_and_legacy_history() { + let thread_id = ThreadId::default(); + + assert_eq!( + resolve_multi_agent_version(&InitialHistory::New, None), + None + ); + assert_eq!( + resolve_multi_agent_version( + &InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Vec::new(), + rollout_path: None, + }), + None, + ), + Some(MultiAgentVersion::V1) + ); + assert_eq!( + resolve_multi_agent_version( + &InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Vec::new(), + rollout_path: None, + }), + Some(MultiAgentVersion::V2), + ), + Some(MultiAgentVersion::V2) + ); + assert_eq!( + resolve_multi_agent_version( + &InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: vec![session_meta_item( + thread_id, + Some(MultiAgentVersion::Disabled) + )], + rollout_path: None, + }), + Some(MultiAgentVersion::V2), + ), + Some(MultiAgentVersion::Disabled) + ); + assert_eq!( + resolve_multi_agent_version(&InitialHistory::Forked(Vec::new()), None), + Some(MultiAgentVersion::V1) + ); +} + #[tokio::test] async fn record_initial_history_new_defers_initial_context_until_first_turn() { let (session, _turn_context) = make_session_and_context().await; @@ -1657,6 +1711,20 @@ async fn record_initial_history_new_defers_initial_context_until_first_turn() { assert_eq!(session.previous_turn_settings().await, None); } +fn session_meta_item( + thread_id: ThreadId, + multi_agent_version: Option, +) -> RolloutItem { + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + multi_agent_version, + ..SessionMeta::default() + }, + git: None, + }) +} + #[tokio::test] async fn resumed_history_injects_initial_context_on_first_context_update_only() { let (session, turn_context) = make_session_and_context().await; @@ -3457,6 +3525,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -4518,6 +4587,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Some(config.multi_agent_version_from_features()), ) .await; @@ -4717,6 +4787,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { &session_telemetry, session_configuration.provider.clone(), &session_configuration, + Some(config.multi_agent_version_from_features()), services.user_shell.as_ref(), services.shell_zsh_path.as_ref(), services.main_execve_wrapper_exe.as_ref(), @@ -4740,6 +4811,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { state: Mutex::new(state), managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), + multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), @@ -4867,6 +4939,7 @@ async fn make_session_with_config_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Some(config.multi_agent_version_from_features()), ) .await?; @@ -4979,6 +5052,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Some(config.multi_agent_version_from_features()), ) .await?; @@ -5959,6 +6033,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: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -6565,6 +6640,7 @@ where &session_telemetry, session_configuration.provider.clone(), &session_configuration, + Some(config.multi_agent_version_from_features()), services.user_shell.as_ref(), services.shell_zsh_path.as_ref(), services.main_execve_wrapper_exe.as_ref(), @@ -6588,6 +6664,7 @@ where state: Mutex::new(state), managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), + multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a0625d65e2..8b1cf15016 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, + inherited_multi_agent_version: None, }) .await .expect("spawn guardian subagent"); @@ -729,6 +730,5 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { }], } ); - drop(codex); } diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 1e0a3f17c5..0ab40b05e4 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -7,7 +7,9 @@ use codex_model_provider::create_model_provider; use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::models::AdditionalPermissionProfile; +use codex_protocol::openai_models::ModelInfo; 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 +80,7 @@ pub struct TurnContext { pub(crate) compact_prompt: Option, pub(crate) user_instructions: Option, pub(crate) collaboration_mode: CollaborationMode, + pub(crate) multi_agent_version: Option, pub(crate) personality: Option, pub(crate) approval_policy: Constrained, pub(crate) permission_profile: PermissionProfile, @@ -101,6 +104,12 @@ pub struct TurnContext { pub(crate) server_model_warning_emitted: AtomicBool, pub(crate) model_verification_emitted: AtomicBool, } + +enum TurnMultiAgentRuntime { + Resolve, + Unset, +} + impl TurnContext { pub(crate) fn permission_profile(&self) -> PermissionProfile { self.permission_profile.clone() @@ -246,6 +255,7 @@ impl TurnContext { compact_prompt: self.compact_prompt.clone(), user_instructions: self.user_instructions.clone(), collaboration_mode, + multi_agent_version: self.multi_agent_version, personality: self.personality, approval_policy: self.approval_policy.clone(), permission_profile: self.permission_profile.clone(), @@ -454,6 +464,7 @@ impl Session { session_telemetry: &SessionTelemetry, provider: ModelProviderInfo, session_configuration: &SessionConfiguration, + multi_agent_version: Option, user_shell: &shell::Shell, shell_zsh_path: Option<&PathBuf>, main_execve_wrapper_exe: Option<&PathBuf>, @@ -543,6 +554,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, personality: session_configuration.personality, approval_policy: session_configuration.approval_policy.clone(), permission_profile: session_configuration.permission_profile(), @@ -684,8 +696,43 @@ impl Session { final_output_json_schema: Option>, turn_environments: ResolvedTurnEnvironments, ) -> Arc { - let primary_turn_environment = turn_environments.primary(); + self.new_turn_context_from_configuration( + sub_id, + session_configuration, + final_output_json_schema, + turn_environments, + TurnMultiAgentRuntime::Resolve, + ) + .await + } + + async fn new_startup_prewarm_turn_from_configuration( + &self, + sub_id: String, + session_configuration: SessionConfiguration, + turn_environments: ResolvedTurnEnvironments, + ) -> Arc { + self.new_turn_context_from_configuration( + sub_id, + session_configuration, + /*final_output_json_schema*/ None, + turn_environments, + TurnMultiAgentRuntime::Unset, + ) + .await + } + + async fn new_turn_context_from_configuration( + &self, + sub_id: String, + session_configuration: SessionConfiguration, + final_output_json_schema: Option>, + turn_environments: ResolvedTurnEnvironments, + multi_agent_runtime: TurnMultiAgentRuntime, + ) -> Arc { + let primary_turn_environment = turn_environments.primary().cloned(); let cwd = primary_turn_environment + .as_ref() .map(|turn_environment| turn_environment.cwd.clone()) .unwrap_or_else(|| session_configuration.cwd.clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); @@ -704,6 +751,14 @@ impl Session { &per_turn_config.to_models_manager_config(), ) .await; + + let multi_agent_version = match multi_agent_runtime { + TurnMultiAgentRuntime::Resolve => Some( + self.resolve_multi_agent_version_for_model(&model_info, &per_turn_config) + .await, + ), + TurnMultiAgentRuntime::Unset => None, + }; let plugin_outcome = self .services .plugins_manager @@ -727,6 +782,7 @@ impl Session { &self.services.session_telemetry, session_configuration.provider.clone(), &session_configuration, + multi_agent_version, self.services.user_shell.as_ref(), self.services.shell_zsh_path.as_ref(), self.services.main_execve_wrapper_exe.as_ref(), @@ -780,6 +836,34 @@ impl Session { } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc { + let (session_configuration, turn_environments) = + self.default_turn_configuration_and_environments().await; + self.new_turn_from_configuration( + sub_id, + session_configuration, + /*final_output_json_schema*/ None, + turn_environments, + ) + .await + } + + pub(crate) async fn new_startup_prewarm_turn_with_sub_id( + &self, + sub_id: String, + ) -> Arc { + let (session_configuration, turn_environments) = + self.default_turn_configuration_and_environments().await; + self.new_startup_prewarm_turn_from_configuration( + sub_id, + session_configuration, + turn_environments, + ) + .await + } + + async fn default_turn_configuration_and_environments( + &self, + ) -> (SessionConfiguration, ResolvedTurnEnvironments) { let session_configuration = { let state = self.state.lock().await; state.session_configuration.clone() @@ -796,14 +880,7 @@ impl Session { ResolvedTurnEnvironments::default() } }; - - self.new_turn_from_configuration( - sub_id, - session_configuration, - /*final_output_json_schema*/ None, - turn_environments, - ) - .await + (session_configuration, turn_environments) } fn overlay_runtime_cwd_on_primary_environment( diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 0669e54469..13e0871e92 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -222,7 +222,7 @@ async fn schedule_startup_prewarm_inner( ) -> CodexResult { let prewarm_started_at = Instant::now(); let startup_turn_context = session - .new_default_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned()) + .new_startup_prewarm_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned()) .await; startup_turn_context.session_telemetry.record_startup_phase( "startup_prewarm_create_turn_context", diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 38a8a96276..a3cd10cf82 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: Option, + ) -> Self { if !config.agent_interrupt_message_enabled { return Self::Disabled; } - if config.features.enabled(Feature::MultiAgentV2) { + if multi_agent_version == Some(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/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 9f890ae0ec..f41c8907f9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -12,6 +12,7 @@ use crate::session::Codex; use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; use crate::session::INITIAL_SUBMIT_ID; +use crate::session::resolve_multi_agent_version; use crate::shell_snapshot::ShellSnapshot; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; @@ -38,6 +39,7 @@ use codex_protocol::openai_models::ModelPreset; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; 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; @@ -191,6 +193,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 @@ -590,14 +593,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 @@ -622,6 +629,7 @@ impl ThreadManager { options.parent_trace, options.environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -648,13 +656,23 @@ 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().ok_or_else(|| { + CodexErr::InvalidRequest("multi-agent version is unresolved".to_string()) + })?; options.initial_history = fork_history_from_snapshot( ForkSnapshot::Interrupted, history, - InterruptedTurnHistoryMarker::from_config(&options.config), + InterruptedTurnHistoryMarker::from_config_and_version( + &options.config, + Some(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( @@ -707,6 +725,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, + /*inherited_multi_agent_version*/ None, )) .await } @@ -734,6 +753,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*user_shell_override*/ Some(user_shell_override), + /*inherited_multi_agent_version*/ None, )) .await } @@ -770,6 +790,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*user_shell_override*/ Some(user_shell_override), + /*inherited_multi_agent_version*/ None, )) .await } @@ -919,7 +940,10 @@ impl ThreadManager { InitialHistory::Forked(_) => history.forked_from_id(), InitialHistory::New | InitialHistory::Cleared => None, }; - let interrupted_marker = InterruptedTurnHistoryMarker::from_config(&config); + let interrupted_marker = InterruptedTurnHistoryMarker::from_config_and_version( + &config, + resolve_multi_agent_version(&history, /*inherited_multi_agent_version*/ None), + ); let history = fork_history_from_snapshot(snapshot, history, interrupted_marker); let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), @@ -939,6 +963,7 @@ impl ThreadManager { parent_trace, environments, /*user_shell_override*/ None, + /*inherited_multi_agent_version*/ None, )) .await } @@ -1062,6 +1087,7 @@ impl ThreadManagerState { /*inherited_shell_snapshot*/ None, /*inherited_exec_policy*/ None, /*environments*/ None, + /*inherited_multi_agent_version*/ None, )) .await } @@ -1080,6 +1106,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) @@ -1101,6 +1128,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -1117,6 +1145,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); @@ -1138,6 +1167,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -1156,6 +1186,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) @@ -1177,6 +1208,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*user_shell_override*/ None, + inherited_multi_agent_version, )) .await } @@ -1198,6 +1230,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, @@ -1216,6 +1249,7 @@ impl ThreadManagerState { parent_trace, environments, user_shell_override, + inherited_multi_agent_version, )) .await } @@ -1239,6 +1273,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 { @@ -1268,6 +1303,22 @@ impl ThreadManagerState { .parent_rollout_thread_trace_for_source(&session_source, &initial_history) .await; let tracked_session_source = session_source.clone(); + let source_thread_id = match &initial_history { + InitialHistory::Resumed(resumed) => Some(resumed.conversation_id), + InitialHistory::Forked(_) => forked_from_thread_id, + InitialHistory::New | InitialHistory::Cleared => None, + }; + let inherited_multi_agent_version = match inherited_multi_agent_version { + Some(multi_agent_version) => Some(multi_agent_version), + None => match source_thread_id { + Some(source_thread_id) => self + .get_thread(source_thread_id) + .await + .ok() + .and_then(|source_thread| source_thread.multi_agent_version()), + None => None, + }, + }; let CodexSpawnOk { codex, thread_id, .. } = Codex::spawn(CodexSpawnArgs { @@ -1298,6 +1349,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), + inherited_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 76880eeb9e..cdc1d90dc0 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1337,7 +1337,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { SnapshotTurnState { ends_mid_turn: true, active_turn_id: Some("turn-explicit".to_string()), - active_turn_start_index: Some(1), + active_turn_start_index: Some(2), }, ); diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 821b3619c4..138a45e137 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -1,6 +1,4 @@ use crate::agent::control::SpawnAgentOptions; -use crate::agent::exceeds_thread_spawn_depth_limit; -use crate::agent::next_thread_spawn_depth; use crate::agent::status::is_final; use crate::config::Config; use crate::function_tool::FunctionCallError; @@ -11,6 +9,7 @@ use crate::tools::handlers::parse_arguments; use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; use codex_protocol::protocol::AgentStatus; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::user_input::UserInput; @@ -89,6 +88,7 @@ struct ReportAgentJobResultToolResult { struct JobRunnerOptions { max_concurrency: usize, spawn_config: Config, + multi_agent_version: MultiAgentVersion, } #[derive(Debug, Clone)] @@ -111,26 +111,32 @@ async fn build_runner_options( turn: &Arc, requested_concurrency: Option, ) -> Result { - let session_source = turn.session_source.clone(); - let child_depth = next_thread_spawn_depth(&session_source); - let max_depth = turn.config.agent_max_depth; - if exceeds_thread_spawn_depth_limit(child_depth, max_depth) { + let multi_agent_version = turn.multi_agent_version.ok_or_else(|| { + FunctionCallError::RespondToModel( + "multi-agent version is unresolved; this session cannot spawn workers".to_string(), + ) + })?; + if multi_agent_version == MultiAgentVersion::Disabled { return Err(FunctionCallError::RespondToModel( - "agent depth limit reached; this session cannot spawn more subagents".to_string(), + "multi-agent runtime is disabled; this session cannot spawn workers".to_string(), )); } - if turn.config.agent_max_threads == Some(0) { + let agent_max_threads = turn + .config + .effective_agent_max_threads(multi_agent_version) + .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; + if agent_max_threads == Some(0) { return Err(FunctionCallError::RespondToModel( "agent thread limit reached; this session cannot spawn more subagents".to_string(), )); } - let max_concurrency = - normalize_concurrency(requested_concurrency, turn.config.agent_max_threads); + let max_concurrency = normalize_concurrency(requested_concurrency, agent_max_threads); let base_instructions = session.get_base_instructions().await; let spawn_config = build_agent_spawn_config(&base_instructions, turn.as_ref())?; Ok(JobRunnerOptions { max_concurrency, spawn_config, + multi_agent_version, }) } @@ -210,6 +216,7 @@ async fn run_agent_job_loop( Some(SessionSource::SubAgent(SubAgentSource::Other(format!( "agent_job:{job_id}" )))), + options.multi_agent_version, SpawnAgentOptions { parent_thread_id: Some(session.conversation_id), environments: Some(turn.environments.to_selections()), 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..c099412987 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,9 @@ 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 multi_agent_version = turn + .multi_agent_version + .ok_or_else(|| FunctionCallError::Fatal("multi-agent version is unresolved".to_string()))?; let result = Box::pin(session.services.agent_control.spawn_agent_with_metadata( config, @@ -121,6 +123,7 @@ async fn handle_spawn_agent( role_name, /*task_name*/ None, )?), + multi_agent_version, SpawnAgentOptions { fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), 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 bfe2ef3d9a..4c33ecad1b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -134,6 +134,11 @@ model_reasoning_effort = "minimal" role_name } +fn set_turn_config(turn: &mut TurnContext, config: crate::config::Config) { + turn.multi_agent_version = Some(config.multi_agent_version_from_features()); + turn.config = Arc::new(config); +} + fn expect_text_output(output: T) -> (String, Option) where T: ToolOutput, @@ -382,6 +387,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: Some(codex_protocol::protocol::MultiAgentVersion::V2), ..turn }; @@ -424,7 +430,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( @@ -891,7 +897,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()) @@ -959,6 +965,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: Some(codex_protocol::protocol::MultiAgentVersion::V2), ..turn }; @@ -1040,7 +1047,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), @@ -1074,7 +1081,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), @@ -1134,7 +1141,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); @@ -1231,7 +1238,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( @@ -1271,7 +1278,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( @@ -1311,7 +1318,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( @@ -1340,18 +1347,22 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { async fn multi_agent_v2_send_message_accepts_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); + let mut config = (*turn.config).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + set_turn_config(&mut turn, config); let root = manager .start_thread((*turn.config).clone()) .await .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - let mut config = (*turn.config).clone(); - config - .features - .enable(Feature::MultiAgentV2) - .expect("test config should allow feature update"); - turn.config = Arc::new(config); + let multi_agent_version = root + .thread + .multi_agent_version() + .expect("root thread should have a multi-agent version"); let child_path = AgentPath::try_from("/root/worker").expect("agent path"); let child_thread_id = session @@ -1371,6 +1382,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { agent_nickname: None, agent_role: None, })), + multi_agent_version, crate::agent::control::SpawnAgentOptions::default(), ) .await @@ -1416,18 +1428,22 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); + let mut config = (*turn.config).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + set_turn_config(&mut turn, config); let root = manager .start_thread((*turn.config).clone()) .await .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - let mut config = (*turn.config).clone(); - config - .features - .enable(Feature::MultiAgentV2) - .expect("test config should allow feature update"); - turn.config = Arc::new(config); + let multi_agent_version = root + .thread + .multi_agent_version() + .expect("root thread should have a multi-agent version"); let child_path = AgentPath::try_from("/root/worker").expect("agent path"); let child_thread_id = session @@ -1447,6 +1463,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { agent_nickname: None, agent_role: None, })), + multi_agent_version, crate::agent::control::SpawnAgentOptions::default(), ) .await @@ -1507,7 +1524,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); @@ -1593,15 +1610,19 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); + let mut config = (*turn.config).clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + set_turn_config(&mut turn, config.clone()); let root = manager .start_thread((*turn.config).clone()) .await .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - let mut config = (*turn.config).clone(); - let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config.clone()); + let multi_agent_version = root + .thread + .multi_agent_version() + .expect("root thread should have a multi-agent version"); 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"); @@ -1622,6 +1643,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { agent_nickname: None, agent_role: None, })), + multi_agent_version, crate::agent::control::SpawnAgentOptions::default(), ) .await @@ -1643,6 +1665,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { agent_nickname: None, agent_role: None, })), + multi_agent_version, crate::agent::control::SpawnAgentOptions::default(), ) .await @@ -1688,7 +1711,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); @@ -1752,7 +1775,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); @@ -1808,7 +1831,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); @@ -1873,15 +1896,15 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); + let mut config = turn.config.as_ref().clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + set_turn_config(&mut turn, config); let root = manager .start_thread((*turn.config).clone()) .await .expect("root thread should start"); session.services.agent_control = manager.agent_control(); session.conversation_id = root.thread_id; - let mut config = turn.config.as_ref().clone(); - let _ = config.features.enable(Feature::MultiAgentV2); - turn.config = Arc::new(config); let session = Arc::new(session); let turn = Arc::new(turn); @@ -2016,7 +2039,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); @@ -2069,7 +2092,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); @@ -2149,7 +2172,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( @@ -2188,7 +2211,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), @@ -2398,7 +2421,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, @@ -2872,7 +2895,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); @@ -2957,7 +2980,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( @@ -2987,7 +3010,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( @@ -3022,7 +3045,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); @@ -3077,7 +3100,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); @@ -3117,7 +3140,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( @@ -3147,7 +3170,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( @@ -3355,7 +3378,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); @@ -3449,7 +3472,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); @@ -3530,7 +3553,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); @@ -3621,7 +3644,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); @@ -3710,7 +3733,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); @@ -3758,7 +3781,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) @@ -3783,7 +3806,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); @@ -3889,7 +3912,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); @@ -4284,7 +4307,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..5d3a5f550c 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,9 @@ 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 multi_agent_version = turn + .multi_agent_version + .ok_or_else(|| FunctionCallError::Fatal("multi-agent version is unresolved".to_string()))?; let spawn_source = thread_spawn_source( session.conversation_id, @@ -141,6 +143,7 @@ async fn handle_spawn_agent( (_, initial_operation) => initial_operation, }, Some(spawn_source), + multi_agent_version, SpawnAgentOptions { fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), fork_mode, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index c7d9ae5aeb..4924b9509c 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; @@ -287,11 +290,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 == Some(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 { + None | Some(MultiAgentVersion::Disabled) => false, + Some(MultiAgentVersion::V1) => !exceeds_thread_spawn_depth_limit( + next_thread_spawn_depth(&turn_context.session_source), + turn_context.config.agent_max_depth, + ), + Some(MultiAgentVersion::V2) => true, + } } fn goal_tools_enabled(turn_context: &TurnContext) -> bool { @@ -303,7 +313,7 @@ fn goal_tools_enabled(turn_context: &TurnContext) -> bool { } fn agent_jobs_tools_enabled(turn_context: &TurnContext) -> bool { - turn_context.features.get().enabled(Feature::SpawnCsv) + turn_context.features.get().enabled(Feature::SpawnCsv) && collab_tools_enabled(turn_context) } 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 e129e4ddaa..dc93846609 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -215,6 +215,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 = Some(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) { diff --git a/codex-rs/core/tests/suite/auto_review.rs b/codex-rs/core/tests/suite/auto_review.rs index b3574cf5e2..570f97008f 100644 --- a/codex-rs/core/tests/suite/auto_review.rs +++ b/codex-rs/core/tests/suite/auto_review.rs @@ -233,6 +233,7 @@ fn remote_model_with_auto_review_override(slug: &str, review_model: &str) -> Mod supports_search_tool: false, auto_review_model_override: Some(review_model.to_string()), tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), diff --git a/codex-rs/core/tests/suite/model_runtime_selectors.rs b/codex-rs/core/tests/suite/model_runtime_selectors.rs index b6f2ea3c7d..a66668e690 100644 --- a/codex-rs/core/tests/suite/model_runtime_selectors.rs +++ b/codex-rs/core/tests/suite/model_runtime_selectors.rs @@ -11,6 +11,7 @@ use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; @@ -31,6 +32,10 @@ use tokio::time::Duration; use tokio::time::Instant; use tokio::time::sleep; +const CHILD_MODEL: &str = "test-multi-agent-child"; +const ROOT_MODEL: &str = "test-multi-agent-root"; +const ROOT_PROMPT: &str = "spawn a child"; + fn remote_model(slug: &str) -> ModelInfo { ModelInfo { visibility: ModelVisibility::List, @@ -171,3 +176,109 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> { Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_multi_agent_selector_overrides_feature_flags() -> Result<()> { + skip_if_no_network!(Ok(())); + + let mut v2_model = remote_model("test-multi-agent-v2"); + v2_model.multi_agent_version = Some(MultiAgentVersion::V2); + let v2_body = response_body_for_remote_model(v2_model, |config| { + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + config + .features + .disable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + }) + .await?; + assert!(tool_names(&v2_body).contains(&"send_message".to_string())); + + let mut disabled_model = remote_model("test-multi-agent-disabled"); + disabled_model.multi_agent_version = Some(MultiAgentVersion::Disabled); + let disabled_body = response_body_for_remote_model(disabled_model, |config| { + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + }) + .await?; + let disabled_tools = tool_names(&disabled_body); + assert!(disabled_tools.iter().all(|name| !matches!( + name.as_str(), + "multi_agent_v1" | "spawn_agent" | "send_message" | "wait_agent" | "list_agents" + ))); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_multi_agent_selector_uses_model_selected_before_first_turn() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = wiremock::MockServer::start().await; + let mut initial_model = remote_model(ROOT_MODEL); + initial_model.multi_agent_version = Some(MultiAgentVersion::V1); + let mut selected_model = remote_model(CHILD_MODEL); + selected_model.multi_agent_version = Some(MultiAgentVersion::V2); + let models_mock = mount_models_once( + &server, + ModelsResponse { + models: vec![initial_model, selected_model], + }, + ) + .await; + let response_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config.model = Some(ROOT_MODEL.to_string()); + }); + let test = builder.build(&server).await?; + assert_eq!( + ( + models_mock.requests().len(), + test.codex.multi_agent_version(), + ), + (1, None) + ); + + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + model: Some(CHILD_MODEL.to_string()), + ..Default::default() + }, + ) + .await?; + assert_eq!(test.codex.multi_agent_version(), None); + + test.submit_turn(ROOT_PROMPT).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + assert_eq!( + ( + models_mock.requests().len(), + test.codex.multi_agent_version(), + tool_names(&response_mock.single_request().body_json()) + .contains(&"send_message".to_string()), + ), + (1, Some(MultiAgentVersion::V2), true) + ); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index f7d2bf578c..b76a3522ec 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -114,6 +114,7 @@ fn test_model_info( supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), @@ -933,6 +934,7 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result< supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), diff --git a/codex-rs/core/tests/suite/models_cache_ttl.rs b/codex-rs/core/tests/suite/models_cache_ttl.rs index 5244a178a4..34fc98b59f 100644 --- a/codex-rs/core/tests/suite/models_cache_ttl.rs +++ b/codex-rs/core/tests/suite/models_cache_ttl.rs @@ -372,5 +372,6 @@ fn test_remote_model(slug: &str, priority: i32) -> ModelInfo { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, } } diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index 4219576d4f..231a6d826f 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -594,6 +594,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, }; let _models_mock = mount_models_once( @@ -706,6 +707,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, }; let _models_mock = mount_models_once( 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/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index 00222a3a3a..fd6e01c075 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -479,6 +479,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), @@ -730,6 +731,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), @@ -1215,6 +1217,7 @@ fn test_remote_model_with_policy( supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index c3d598c7ba..36ea7d4b2c 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -1351,6 +1351,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, }], }, ) diff --git a/codex-rs/core/tests/suite/spawn_agent_description.rs b/codex-rs/core/tests/suite/spawn_agent_description.rs index 787ba10501..c1f0d522c7 100644 --- a/codex-rs/core/tests/suite/spawn_agent_description.rs +++ b/codex-rs/core/tests/suite/spawn_agent_description.rs @@ -62,6 +62,7 @@ fn test_model_info( supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers, 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/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 0c6cdf6d8b..2e34475387 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -1357,6 +1357,7 @@ async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> an supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, priority: 1, additional_speed_tiers: Vec::new(), service_tiers: Vec::new(), diff --git a/codex-rs/models-manager/src/model_info.rs b/codex-rs/models-manager/src/model_info.rs index 379d173e4e..58137a3f50 100644 --- a/codex-rs/models-manager/src/model_info.rs +++ b/codex-rs/models-manager/src/model_info.rs @@ -101,6 +101,7 @@ pub fn model_info_from_slug(slug: &str) -> ModelInfo { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, } } diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 0214f450fd..4b678a9209 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -22,6 +22,7 @@ use crate::config_types::ReasoningSummary; use crate::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use crate::config_types::ServiceTier; use crate::config_types::Verbosity; +use crate::protocol::MultiAgentVersion; const PERSONALITY_PLACEHOLDER: &str = "{{ personality }}"; pub const SPEED_TIER_FAST: &str = "fast"; @@ -347,6 +348,12 @@ pub struct ModelInfo { deserialize_with = "deserialize_optional_model_selector" )] pub tool_mode: Option, + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_optional_model_selector" + )] + pub multi_agent_version: Option, } impl ModelInfo { @@ -643,6 +650,7 @@ mod tests { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, } } @@ -901,6 +909,22 @@ mod tests { assert!(!object.contains_key("tool_mode")); } + #[test] + fn model_info_treats_unknown_multi_agent_version_as_omitted() { + let mut value = + serde_json::to_value(test_model(/*spec*/ None)).expect("serialize test model"); + let object = value + .as_object_mut() + .expect("model info should be an object"); + object.insert( + "multi_agent_version".to_string(), + serde_json::Value::String("future_multi_agent_version".to_string()), + ); + let model = serde_json::from_value::(value).expect("deserialize model info"); + + assert_eq!(model.multi_agent_version, None); + } + #[test] fn resolved_context_window_prefers_context_window() { let model = ModelInfo { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 7c91223f72..4619ed3eb5 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2441,6 +2441,18 @@ 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, /*thread_id*/ 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)) @@ -2711,6 +2723,32 @@ 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) => + { + Some(meta_line) + } + _ => None, + }) + .and_then(|meta_line| meta_line.meta.multi_agent_version) +} + +#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +pub enum MultiAgentVersion { + Disabled, + 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 @@ -2750,6 +2788,8 @@ pub struct SessionMeta { pub dynamic_tools: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub memory_mode: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub multi_agent_version: Option, } impl Default for SessionMeta { @@ -2771,6 +2811,7 @@ impl Default for SessionMeta { base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, } } } diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 8eb2380536..caad965223 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -447,6 +447,7 @@ fn write_rollout(path: &std::path::Path, thread_id: ThreadId, message: &str) -> base_instructions: None, dynamic_tools: None, memory_mode: None, + multi_agent_version: None, }, git: None, }; 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..d102d5c540 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: Option, }, Resume { path: PathBuf, @@ -172,9 +174,24 @@ impl RolloutRecorderParams { thread_source, base_instructions, dynamic_tools, + multi_agent_version: None, } } + pub fn with_multi_agent_version( + mut self, + multi_agent_version: Option, + ) -> Self { + if let Self::Create { + multi_agent_version: version, + .. + } = &mut self + { + *version = multi_agent_version; + } + self + } + pub fn resume(path: PathBuf) -> Self { Self::Resume { path } } @@ -661,6 +678,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 +714,7 @@ impl RolloutRecorder { Some(dynamic_tools) }, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), + 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..4da081eff5 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, }; 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/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 22a3cb803a..a84ede199b 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -158,6 +158,7 @@ fn write_rollout_with_user_message( 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..4ae629d79c 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; @@ -164,7 +169,31 @@ impl ThreadStore for InMemoryThreadStore { async fn create_thread(&self, params: CreateThreadParams) -> 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: 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(()) } @@ -172,13 +201,50 @@ impl ThreadStore for InMemoryThreadStore { async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> { let mut state = self.state.lock().await; state.calls.resume_thread += 1; - state.histories.entry(params.thread_id).or_default(); + if let Some(history) = params.history { + state.histories.insert(params.thread_id, history); + } else { + state.histories.entry(params.thread_id).or_default(); + } if let Some(rollout_path) = params.rollout_path { state.rollout_paths.insert(rollout_path, params.thread_id); } Ok(()) } + async fn set_multi_agent_version_if_unset( + &self, + params: SetMultiAgentVersionIfUnsetParams, + ) -> ThreadStoreResult { + let mut state = self.state.lock().await; + let history = + state + .histories + .get_mut(¶ms.thread_id) + .ok_or(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + })?; + let mut session_meta = history + .iter() + .rev() + .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), + })?; + if let Some(multi_agent_version) = session_meta.meta.multi_agent_version { + return Ok(multi_agent_version); + } + 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..8c630f9956 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -35,7 +35,8 @@ pub(super) async fn create_thread( params.thread_source, params.base_instructions, params.dynamic_tools, - ), + ) + .with_multi_agent_version(params.multi_agent_version), ) .await .map_err(|err| ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index e26c314e73..234c98dcd2 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_if_unset; mod unarchive_thread; mod update_thread_metadata; @@ -20,6 +21,7 @@ use std::collections::hash_map::Entry; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::Mutex; +use tokio::sync::Semaphore; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; @@ -30,6 +32,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 +41,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 +60,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, + multi_agent_version_update_gate: Arc, state_db: Option, } @@ -95,6 +100,7 @@ impl LocalThreadStore { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), + multi_agent_version_update_gate: Arc::new(Semaphore::new(1)), state_db, } } @@ -180,6 +186,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_if_unset::set_multi_agent_version_if_unset(self, params).await + } + async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> { live_writer::append_items(self, params).await } @@ -760,6 +773,44 @@ mod tests { assert!(err.to_string().contains("already has a live local writer")); } + #[tokio::test] + async fn set_multi_agent_version_if_unset_appends_once_and_returns_existing_value() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let thread_id = ThreadId::default(); + + store + .create_thread(create_thread_params(thread_id)) + .await + .expect("create thread"); + + let first = store + .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V2, + }) + .await + .expect("set missing runtime"); + assert_eq!(first, MultiAgentVersion::V2); + assert_eq!( + session_meta_versions(&store, thread_id).await, + vec![None, Some(MultiAgentVersion::V2)] + ); + + let second = store + .set_multi_agent_version_if_unset(SetMultiAgentVersionIfUnsetParams { + thread_id, + multi_agent_version: MultiAgentVersion::V1, + }) + .await + .expect("return existing runtime"); + assert_eq!(second, MultiAgentVersion::V2); + assert_eq!( + session_meta_versions(&store, thread_id).await, + vec![None, Some(MultiAgentVersion::V2)] + ); + } + #[tokio::test] async fn resume_thread_rejects_duplicate_live_writer() { let home = TempDir::new().expect("temp dir"); @@ -1028,6 +1079,7 @@ mod tests { thread_source: None, base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + multi_agent_version: None, metadata: thread_metadata(), event_persistence_mode: ThreadEventPersistenceMode::Limited, } @@ -1052,6 +1104,28 @@ mod tests { })) } + async fn session_meta_versions( + store: &LocalThreadStore, + thread_id: ThreadId, + ) -> Vec> { + store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: true, + }) + .await + .expect("load history") + .items + .into_iter() + .filter_map(|item| match item { + RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == thread_id => { + Some(meta_line.meta.multi_agent_version) + } + _ => None, + }) + .collect() + } + async fn assert_rollout_contains_message(path: &std::path::Path, expected: &str) { let (items, _, _) = RolloutRecorder::load_rollout_items(path) .await diff --git a/codex-rs/thread-store/src/local/set_multi_agent_version_if_unset.rs b/codex-rs/thread-store/src/local/set_multi_agent_version_if_unset.rs new file mode 100644 index 0000000000..aadc5e3f69 --- /dev/null +++ b/codex-rs/thread-store/src/local/set_multi_agent_version_if_unset.rs @@ -0,0 +1,88 @@ +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::RolloutItem; +use codex_rollout::append_rollout_item_to_path; + +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 _permit = store + .multi_agent_version_update_gate + .acquire() + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to set thread multi-agent version: {err}"), + })?; + 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: true, + 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), + })?; + let mut session_meta = history + .items + .iter() + .rev() + .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), + })?; + if let Some(multi_agent_version) = session_meta.meta.multi_agent_version { + return Ok(multi_agent_version); + } + + 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..ef6fc18285 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<()>; + /// Sets a thread's multi-agent runtime if it is unset, returning the stored value. + 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..1f4fc39a01 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,15 @@ pub struct ThreadPersistenceMetadata { pub memory_mode: MemoryMode, } +/// Parameters for setting a thread's multi-agent runtime when it has not been recorded yet. +#[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, +} + /// Parameters required to create a persisted thread. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct CreateThreadParams { @@ -82,6 +92,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: Option, /// Metadata captured for the newly created thread. pub metadata: ThreadPersistenceMetadata, /// Whether persistence should include the extended event surface. diff --git a/codex-rs/tools/src/tool_config_tests.rs b/codex-rs/tools/src/tool_config_tests.rs index 165d6efa20..1b4d789394 100644 --- a/codex-rs/tools/src/tool_config_tests.rs +++ b/codex-rs/tools/src/tool_config_tests.rs @@ -46,6 +46,7 @@ fn model_with_shell_type(shell_type: ConfigShellToolType) -> ModelInfo { supports_search_tool: false, auto_review_model_override: None, tool_mode: None, + multi_agent_version: None, } }