diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 7cf5d00b4..8df3b0a07 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -109,6 +109,79 @@ impl AgentControl { .await } + pub(crate) async fn ensure_v2_agent_loaded( + &self, + config: Config, + thread_id: ThreadId, + ) -> CodexResult<()> { + let state = self.upgrade()?; + if state.get_thread(thread_id).await.is_ok() { + return Ok(()); + } + if self.state.agent_metadata_for_thread(thread_id).is_none() { + return Err(CodexErr::ThreadNotFound(thread_id)); + } + + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: true, + }) + .await?; + let stored_source = stored_thread.source.clone(); + let stored_parent_thread_id = stored_thread.parent_thread_id; + let history = stored_thread + .history + .ok_or(CodexErr::ThreadNotFound(thread_id))? + .items; + let initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history, + rollout_path: stored_thread.rollout_path, + }); + if initial_history.get_multi_agent_version() != Some(MultiAgentVersion::V2) { + return Err(CodexErr::ThreadNotFound(thread_id)); + } + + let (session_source, _) = initial_history + .get_resumed_session_sources() + .unwrap_or((stored_source, None)); + let parent_thread_id = initial_history + .get_resumed_parent_thread_id() + .or(stored_parent_thread_id); + let inherited_shell_snapshot = self + .inherited_shell_snapshot_for_source(&state, Some(&session_source)) + .await; + let inherited_exec_policy = self + .inherited_exec_policy_for_source(&state, Some(&session_source), &config) + .await; + + match state + .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { + config, + initial_history, + agent_control: self.clone(), + session_source, + parent_thread_id, + inherited_shell_snapshot, + inherited_exec_policy, + }) + .await + { + Ok(reloaded_thread) => { + state.notify_thread_created(reloaded_thread.thread_id); + Ok(()) + } + Err(err) => { + if state.get_thread(thread_id).await.is_ok() { + return Ok(()); + } + Err(err) + } + } + } + async fn spawn_agent_internal( &self, config: Config, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3866d2201..c9cd44492 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -527,6 +527,113 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig )); } +#[tokio::test] +async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() { + let (home, mut config) = test_config().await; + let _ = config.features.enable(Feature::MultiAgentV2); + let _ = config.features.enable(Feature::Sqlite); + let state_db = init_state_db(&config).await; + let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + state_db.clone(), + ); + let control = manager.agent_control(); + let harness = AgentControlHarness { + _home: home, + config, + state_db, + manager, + control, + }; + let (parent_thread_id, _parent_thread) = harness.start_thread().await; + let agent_path = AgentPath::try_from("/root/worker").expect("agent path"); + let spawned_agent = harness + .control + .spawn_agent_with_metadata( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: Some(agent_path.clone()), + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + parent_thread_id: Some(parent_thread_id), + ..Default::default() + }, + ) + .await + .expect("spawn_agent should succeed"); + let child_thread = harness + .manager + .get_thread(spawned_agent.thread_id) + .await + .expect("child thread should exist"); + child_thread + .inject_response_items(vec![assistant_message( + "child persisted", + Some(MessagePhase::FinalAnswer), + )]) + .await + .expect("child rollout should persist with v2 metadata"); + child_thread + .shutdown_and_wait() + .await + .expect("child thread should shut down"); + + assert!( + harness + .manager + .remove_thread(&spawned_agent.thread_id) + .await + .is_some() + ); + match harness.manager.get_thread(spawned_agent.thread_id).await { + Err(CodexErr::ThreadNotFound(id)) => assert_eq!(id, spawned_agent.thread_id), + Err(err) => panic!("expected ThreadNotFound, got {err:?}"), + Ok(_) => panic!("expected thread to be removed"), + } + + harness + .control + .ensure_v2_agent_loaded(harness.config.clone(), spawned_agent.thread_id) + .await + .expect("known v2 agent should reload"); + let _ = harness + .manager + .get_thread(spawned_agent.thread_id) + .await + .expect("reloaded child thread should exist"); + + let communication = InterAgentCommunication::new( + AgentPath::root(), + agent_path, + Vec::new(), + "hello after reload".to_string(), + /*trigger_turn*/ false, + ); + harness + .control + .send_inter_agent_communication(spawned_agent.thread_id, communication.clone()) + .await + .expect("send_inter_agent_communication should succeed after reload"); + let expected = ( + spawned_agent.thread_id, + Op::InterAgentCommunication { communication }, + ); + let captured = harness + .manager + .captured_ops() + .into_iter() + .find(|entry| *entry == expected); + assert_eq!(captured, Some(expected)); +} + #[tokio::test] async fn encrypted_inter_agent_communication_clears_existing_last_task_message() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index 1b7018939..6e28d0a1d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -42,8 +42,17 @@ impl ToolExecutor for Handler { let receiver_agent = session .services .agent_control - .get_agent_metadata(receiver_thread_id) - .unwrap_or_default(); + .get_agent_metadata(receiver_thread_id); + if receiver_agent.is_some() { + let resume_config = build_agent_resume_config(turn.as_ref())?; + session + .services + .agent_control + .ensure_v2_agent_loaded(resume_config, receiver_thread_id) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err))?; + } + let receiver_agent = receiver_agent.unwrap_or_default(); if args.interrupt { session .services diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs index f23d700e9..528fa5526 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs @@ -75,7 +75,11 @@ pub(crate) async fn handle_message_string_tool( .services .agent_control .get_agent_metadata(receiver_thread_id) - .unwrap_or_default(); + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "agent with id {receiver_thread_id} not found" + )) + })?; if mode == MessageDeliveryMode::TriggerTurn && receiver_agent .agent_path @@ -86,6 +90,16 @@ pub(crate) async fn handle_message_string_tool( "Follow-up tasks can't target the root agent".to_string(), )); } + let receiver_agent_path = receiver_agent.agent_path.clone().ok_or_else(|| { + FunctionCallError::RespondToModel("target agent is missing an agent_path".to_string()) + })?; + let resume_config = build_agent_resume_config(turn.as_ref())?; + session + .services + .agent_control + .ensure_v2_agent_loaded(resume_config, receiver_thread_id) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err))?; session .send_event( &turn, @@ -99,9 +113,6 @@ pub(crate) async fn handle_message_string_tool( .into(), ) .await; - let receiver_agent_path = receiver_agent.agent_path.clone().ok_or_else(|| { - FunctionCallError::RespondToModel("target agent is missing an agent_path".to_string()) - })?; let author = turn .session_source .get_agent_path()