diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 85e54e2fb9..0c8a4d0d96 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -540,6 +540,32 @@ impl AgentControl { }); } + fn prepare_agent_metadata( + &self, + reservation: &mut crate::agent::registry::SpawnReservation, + config: &Config, + agent_path: Option, + agent_role: Option, + preferred_agent_nickname: Option, + ) -> CodexResult { + if let Some(agent_path) = agent_path.as_ref() { + reservation.reserve_agent_path(agent_path)?; + } + let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref()); + let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect(); + let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference( + &candidate_name_refs, + preferred_agent_nickname.as_deref(), + )?); + Ok(AgentMetadata { + agent_id: None, + agent_path, + agent_nickname, + agent_role, + last_task_message: None, + }) + } + #[allow(clippy::too_many_arguments)] fn prepare_thread_spawn( &self, @@ -554,29 +580,20 @@ impl AgentControl { if depth == 1 { self.state.register_root_thread(parent_thread_id); } - if let Some(agent_path) = agent_path.as_ref() { - reservation.reserve_agent_path(agent_path)?; - } - let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref()); - let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect(); - let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference( - &candidate_name_refs, - preferred_agent_nickname.as_deref(), - )?); + let agent_metadata = self.prepare_agent_metadata( + reservation, + config, + agent_path, + agent_role, + preferred_agent_nickname, + )?; let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, depth, - agent_path: agent_path.clone(), - agent_nickname: agent_nickname.clone(), - agent_role: agent_role.clone(), + agent_path: agent_metadata.agent_path.clone(), + agent_nickname: agent_metadata.agent_nickname.clone(), + agent_role: agent_metadata.agent_role.clone(), }); - let agent_metadata = AgentMetadata { - agent_id: None, - agent_path, - agent_nickname, - agent_role, - last_task_message: None, - }; Ok((session_source, agent_metadata)) } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 9078434f1f..649c104f30 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -95,6 +95,77 @@ fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: & } impl AgentControl { + /// Restore persisted V2 agent identities without reopening their runtimes. + pub(crate) async fn restore_v2_agent_metadata( + &self, + config: &Config, + root_thread_id: ThreadId, + ) { + self.state.register_root_thread(root_thread_id); + + let Ok(state) = self.upgrade() else { + return; + }; + let Some(agent_graph_store) = state.agent_graph_store() else { + return; + }; + let descendant_ids = match agent_graph_store + .list_thread_spawn_descendants( + root_thread_id, + Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open), + ) + .await + { + Ok(descendant_ids) => descendant_ids, + Err(err) => { + warn!("failed to restore persisted V2 agent metadata for {root_thread_id}: {err}"); + return; + } + }; + + for thread_id in descendant_ids { + if self.state.agent_metadata_for_thread(thread_id).is_some() { + continue; + } + let restore_result = async { + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: false, + }) + .await?; + let stored_agent_path = stored_thread + .agent_path + .as_deref() + .map(AgentPath::try_from) + .transpose() + .map_err(|err| { + CodexErr::InvalidRequest(format!("invalid stored agent path: {err}")) + })?; + let mut reservation = self.state.reserve_spawn_slot(/*max_threads*/ None)?; + let mut metadata = self.prepare_agent_metadata( + &mut reservation, + config, + stored_agent_path.or_else(|| stored_thread.source.get_agent_path()), + stored_thread + .agent_role + .or_else(|| stored_thread.source.get_agent_role()), + stored_thread + .agent_nickname + .or_else(|| stored_thread.source.get_nickname()), + )?; + metadata.agent_id = Some(thread_id); + reservation.commit(metadata); + Ok::<(), CodexErr>(()) + } + .await; + if let Err(err) = restore_result { + warn!("failed to restore V2 agent metadata for {thread_id}: {err}"); + } + } + } + /// Spawn a new agent thread and submit the initial prompt. #[cfg(test)] pub(crate) async fn spawn_agent( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 66fbe42984..9775f25a7c 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -793,6 +793,14 @@ impl ThreadManager { let (session_source, thread_source) = initial_history .get_resumed_session_sources() .unwrap_or_else(|| (self.state.session_source.clone(), None)); + if let InitialHistory::Resumed(resumed) = &initial_history + && initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2) + && !session_source.is_non_root_agent() + { + agent_control + .restore_v2_agent_metadata(&config, resumed.conversation_id) + .await; + } Box::pin(self.state.spawn_thread_with_source( config, initial_history, diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index d224c1fde6..2df01f543d 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -80,6 +80,7 @@ mod model_visible_layout; mod models_cache_ttl; mod models_etag_responses; mod multi_agent_mode; +mod multi_agent_resume; mod network_approval; mod openai_file_mcp; mod otel; diff --git a/codex-rs/core/tests/suite/multi_agent_resume.rs b/codex-rs/core/tests/suite/multi_agent_resume.rs new file mode 100644 index 0000000000..cbb7dd42b1 --- /dev/null +++ b/codex-rs/core/tests/suite/multi_agent_resume.rs @@ -0,0 +1,247 @@ +use anyhow::Result; +use codex_features::Feature; +use codex_protocol::protocol::AgentStatus; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call_with_namespace; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once_match; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::time::Duration; +use tokio::time::Instant; +use tokio::time::sleep; + +const COLLABORATION_NAMESPACE: &str = "collaboration"; +const SPAWN_CALL_ID: &str = "spawn-worker"; +const FOLLOWUP_CALL_ID: &str = "followup-worker"; +const INITIAL_PROMPT: &str = "spawn a durable worker"; +const INITIAL_TASK: &str = "inspect the repository"; +const FOLLOWUP_PROMPT: &str = "continue the durable worker"; +const FOLLOWUP_TASK: &str = "inspect the tests too"; + +fn decoded_body(request: &wiremock::Request) -> Option> { + let is_zstd = request + .headers + .get("content-encoding") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(',') + .any(|entry| entry.trim().eq_ignore_ascii_case("zstd")) + }); + if is_zstd { + zstd::stream::decode_all(std::io::Cursor::new(&request.body)).ok() + } else { + Some(request.body.clone()) + } +} + +fn body_contains(request: &wiremock::Request, text: &str) -> bool { + decoded_body(request) + .and_then(|body| String::from_utf8(body).ok()) + .is_some_and(|body| body.contains(text)) +} + +fn request_has_input_type(request: &wiremock::Request, input_type: &str) -> bool { + decoded_body(request) + .and_then(|body| serde_json::from_slice::(&body).ok()) + .and_then(|body| body.get("input").and_then(Value::as_array).cloned()) + .is_some_and(|items| { + items + .iter() + .any(|item| item.get("type").and_then(Value::as_str) == Some(input_type)) + }) +} + +fn configure_multi_agent_v2(config: &mut codex_core::config::Config) { + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup() -> Result<()> { + let server = start_mock_server().await; + let spawn_args = serde_json::to_string(&json!({ + "message": INITIAL_TASK, + "task_name": "worker", + }))?; + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, INITIAL_PROMPT), + sse(vec![ + ev_response_created("resp-spawn-1"), + ev_function_call_with_namespace( + SPAWN_CALL_ID, + COLLABORATION_NAMESPACE, + "spawn_agent", + &spawn_args, + ), + ev_completed("resp-spawn-1"), + ]), + ) + .await; + let initial_child_request = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + request_has_input_type(request, "agent_message") && body_contains(request, INITIAL_TASK) + }, + sse(vec![ + ev_response_created("resp-worker-1"), + ev_assistant_message("msg-worker-1", "initial task complete"), + ev_completed("resp-worker-1"), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, SPAWN_CALL_ID) + && !request_has_input_type(request, "agent_message") + }, + sse(vec![ + ev_response_created("resp-spawn-2"), + ev_assistant_message("msg-spawn-2", "worker spawned"), + ev_completed("resp-spawn-2"), + ]), + ) + .await; + + let mut initial_builder = test_codex().with_config(configure_multi_agent_v2); + let initial = initial_builder.build(&server).await?; + let root_thread_id = initial.session_configured.thread_id; + let home = initial.home.clone(); + let rollout_path = initial + .codex + .rollout_path() + .expect("root rollout path") + .to_path_buf(); + initial.submit_turn(INITIAL_PROMPT).await?; + + let deadline = Instant::now() + Duration::from_secs(2); + let worker_thread_id = loop { + if let Some(thread_id) = initial + .thread_manager + .list_thread_ids() + .await + .into_iter() + .find(|thread_id| *thread_id != root_thread_id) + { + break thread_id; + } + if Instant::now() >= deadline { + anyhow::bail!("timed out waiting for spawned worker"); + } + sleep(Duration::from_millis(10)).await; + }; + let worker_thread = initial.thread_manager.get_thread(worker_thread_id).await?; + let deadline = Instant::now() + Duration::from_secs(2); + loop { + if matches!( + worker_thread.agent_status().await, + AgentStatus::Completed(_) + ) { + break; + } + if Instant::now() >= deadline { + anyhow::bail!("timed out waiting for worker completion"); + } + sleep(Duration::from_millis(10)).await; + } + assert!( + initial_child_request + .requests() + .iter() + .any(|request| request.body_contains_text(INITIAL_TASK)) + ); + worker_thread.flush_rollout().await?; + initial.codex.flush_rollout().await?; + drop(worker_thread); + drop(initial); + + let followup_args = serde_json::to_string(&json!({ + "target": "worker", + "message": FOLLOWUP_TASK, + }))?; + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, FOLLOWUP_PROMPT), + sse(vec![ + ev_response_created("resp-followup-1"), + ev_function_call_with_namespace( + FOLLOWUP_CALL_ID, + COLLABORATION_NAMESPACE, + "followup_task", + &followup_args, + ), + ev_completed("resp-followup-1"), + ]), + ) + .await; + let followup_child_request = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + request_has_input_type(request, "agent_message") + && body_contains(request, FOLLOWUP_TASK) + }, + sse(vec![ + ev_response_created("resp-worker-2"), + ev_assistant_message("msg-worker-2", "follow-up complete"), + ev_completed("resp-worker-2"), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, FOLLOWUP_CALL_ID) + && !request_has_input_type(request, "agent_message") + }, + sse(vec![ + ev_response_created("resp-followup-2"), + ev_assistant_message("msg-followup-2", "follow-up sent"), + ev_completed("resp-followup-2"), + ]), + ) + .await; + + let mut resume_builder = test_codex().with_config(configure_multi_agent_v2); + let resumed = resume_builder.resume(&server, home, rollout_path).await?; + assert_eq!( + resumed.thread_manager.list_thread_ids().await, + vec![root_thread_id] + ); + assert!( + resumed + .thread_manager + .get_thread(worker_thread_id) + .await + .is_err() + ); + + resumed.submit_turn(FOLLOWUP_PROMPT).await?; + + assert!( + followup_child_request + .requests() + .iter() + .any(|request| request.body_contains_text(FOLLOWUP_TASK)) + ); + resumed + .thread_manager + .get_thread(worker_thread_id) + .await + .expect("follow-up should lazily reload the original worker"); + + Ok(()) +}