From d3603ae5d38ab3addbf995ee8c51a22ceb068872 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 26 Feb 2026 18:01:53 +0000 Subject: [PATCH] feat: fork thread multi agent (#12499) --- codex-rs/core/src/agent/control.rs | 345 +++++++++++++++++- codex-rs/core/src/thread_manager.rs | 21 ++ .../core/src/tools/handlers/multi_agents.rs | 8 +- codex-rs/core/src/tools/spec.rs | 9 + .../tests/suite/subagent_notifications.rs | 115 ++++++ 5 files changed, 494 insertions(+), 4 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index f805a9e74..bb1488b76 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -4,11 +4,16 @@ use crate::agent::status::is_final; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::find_thread_path_by_id_str; +use crate::rollout::RolloutRecorder; use crate::session_prefix::format_subagent_notification_message; use crate::state_db; use crate::thread_manager::ThreadManagerState; use codex_protocol::ThreadId; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TokenUsage; @@ -18,6 +23,12 @@ use std::sync::Weak; use tokio::sync::watch; const AGENT_NAMES: &str = include_str!("agent_names.txt"); +const FORKED_SPAWN_AGENT_OUTPUT_MESSAGE: &str = "You are the newly spawned agent. The prior conversation history was forked from your parent agent. Treat the next user message as your new task, and use the forked history only as background context."; + +#[derive(Clone, Debug, Default)] +pub(crate) struct SpawnAgentOptions { + pub(crate) fork_parent_spawn_call_id: Option, +} fn agent_nickname_list() -> Vec<&'static str> { AGENT_NAMES @@ -57,6 +68,17 @@ impl AgentControl { config: crate::config::Config, items: Vec, session_source: Option, + ) -> CodexResult { + self.spawn_agent_with_options(config, items, session_source, SpawnAgentOptions::default()) + .await + } + + pub(crate) async fn spawn_agent_with_options( + &self, + config: crate::config::Config, + items: Vec, + session_source: Option, + options: SpawnAgentOptions, ) -> CodexResult { let state = self.upgrade()?; let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; @@ -82,9 +104,75 @@ impl AgentControl { // The same `AgentControl` is sent to spawn the thread. let new_thread = match session_source { Some(session_source) => { - state - .spawn_new_thread_with_source(config, self.clone(), session_source, false, None) - .await? + if let Some(call_id) = options.fork_parent_spawn_call_id.as_ref() { + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + .. + }) = session_source.clone() + else { + return Err(CodexErr::Fatal( + "spawn_agent fork requires a thread-spawn session source".to_string(), + )); + }; + let parent_thread = state.get_thread(parent_thread_id).await.ok(); + if let Some(parent_thread) = parent_thread.as_ref() { + // `record_conversation_items` only queues rollout writes asynchronously. + // Flush/materialize the live parent before snapshotting JSONL for a fork. + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread.codex.session.flush_rollout().await; + } + let rollout_path = parent_thread + .as_ref() + .and_then(|parent_thread| parent_thread.rollout_path()) + .or(find_thread_path_by_id_str( + config.codex_home.as_path(), + &parent_thread_id.to_string(), + ) + .await?) + .ok_or_else(|| { + CodexErr::Fatal(format!( + "parent thread rollout unavailable for fork: {parent_thread_id}" + )) + })?; + let mut forked_rollout_items = + RolloutRecorder::get_rollout_history(&rollout_path) + .await? + .get_rollout_items(); + let mut output = FunctionCallOutputPayload::from_text( + FORKED_SPAWN_AGENT_OUTPUT_MESSAGE.to_string(), + ); + output.success = Some(true); + forked_rollout_items.push(RolloutItem::ResponseItem( + ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output, + }, + )); + let initial_history = InitialHistory::Forked(forked_rollout_items); + state + .fork_thread_with_source( + config, + initial_history, + self.clone(), + session_source, + false, + ) + .await? + } else { + state + .spawn_new_thread_with_source( + config, + self.clone(), + session_source, + false, + None, + ) + .await? + } } None => state.spawn_new_thread(config, self.clone()).await?, }; @@ -421,6 +509,21 @@ mod tests { }) } + /// Returns true when any message item contains `needle` in a text span. + fn history_contains_text(history_items: &[ResponseItem], needle: &str) -> bool { + history_items.iter().any(|item| { + let ResponseItem::Message { content, .. } = item else { + return false; + }; + content.iter().any(|content_item| match content_item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + text.contains(needle) + } + ContentItem::InputImage { .. } => false, + }) + }) + } + async fn wait_for_subagent_notification(parent_thread: &Arc) -> bool { let wait = async { loop { @@ -673,6 +776,242 @@ mod tests { assert_eq!(captured, Some(expected)); } + #[tokio::test] + async fn spawn_agent_can_fork_parent_thread_history() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + parent_thread + .inject_user_message_without_turn("parent seed context".to_string()) + .await; + let turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-history".to_string(); + let parent_spawn_call = ResponseItem::FunctionCall { + id: None, + name: "spawn_agent".to_string(), + arguments: "{}".to_string(), + call_id: parent_spawn_call_id.clone(), + }; + parent_thread + .codex + .session + .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread.codex.session.flush_rollout().await; + + let child_thread_id = harness + .control + .spawn_agent_with_options( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id), + }, + ) + .await + .expect("forked spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + assert_ne!(child_thread_id, parent_thread_id); + let history = child_thread.codex.session.clone_history().await; + assert!(history_contains_text( + history.raw_items(), + "parent seed context" + )); + + let expected = ( + child_thread_id, + Op::UserInput { + items: vec![UserInput::Text { + text: "child task".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }, + ); + let captured = harness + .manager + .captured_ops() + .into_iter() + .find(|entry| *entry == expected); + assert_eq!(captured, Some(expected)); + + let _ = harness + .control + .shutdown_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); + } + + #[tokio::test] + async fn spawn_agent_fork_injects_output_for_parent_spawn_call() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + let turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-1".to_string(); + let parent_spawn_call = ResponseItem::FunctionCall { + id: None, + name: "spawn_agent".to_string(), + arguments: "{}".to_string(), + call_id: parent_spawn_call_id.clone(), + }; + parent_thread + .codex + .session + .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread.codex.session.flush_rollout().await; + + let child_thread_id = harness + .control + .spawn_agent_with_options( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), + }, + ) + .await + .expect("forked spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let history = child_thread.codex.session.clone_history().await; + let injected_output = history.raw_items().iter().find_map(|item| match item { + ResponseItem::FunctionCallOutput { call_id, output } + if call_id == &parent_spawn_call_id => + { + Some(output) + } + _ => None, + }); + let injected_output = + injected_output.expect("forked child should contain synthetic tool output"); + assert_eq!( + injected_output.text_content(), + Some(FORKED_SPAWN_AGENT_OUTPUT_MESSAGE) + ); + assert_eq!(injected_output.success, Some(true)); + + let _ = harness + .control + .shutdown_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); + } + + #[tokio::test] + async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + let turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-unflushed".to_string(); + let parent_spawn_call = ResponseItem::FunctionCall { + id: None, + name: "spawn_agent".to_string(), + arguments: "{}".to_string(), + call_id: parent_spawn_call_id.clone(), + }; + parent_thread + .codex + .session + .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) + .await; + + let child_thread_id = harness + .control + .spawn_agent_with_options( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), + }, + ) + .await + .expect("forked spawn should flush parent rollout before loading history"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let history = child_thread.codex.session.clone_history().await; + + let mut parent_call_index = None; + let mut injected_output_index = None; + for (idx, item) in history.raw_items().iter().enumerate() { + match item { + ResponseItem::FunctionCall { call_id, .. } if call_id == &parent_spawn_call_id => { + parent_call_index = Some(idx); + } + ResponseItem::FunctionCallOutput { call_id, .. } + if call_id == &parent_spawn_call_id => + { + injected_output_index = Some(idx); + } + _ => {} + } + } + + let parent_call_index = + parent_call_index.expect("forked child should include the parent spawn_agent call"); + let injected_output_index = injected_output_index + .expect("forked child should include synthetic output for the parent spawn_agent call"); + assert!(parent_call_index < injected_output_index); + + let _ = harness + .control + .shutdown_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); + } + #[tokio::test] async fn spawn_agent_respects_max_threads_limit() { let max_threads = 1usize; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index f653b3dfc..ab2e4c2d2 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -495,6 +495,27 @@ impl ThreadManagerState { .await } + pub(crate) async fn fork_thread_with_source( + &self, + config: Config, + initial_history: InitialHistory, + agent_control: AgentControl, + session_source: SessionSource, + persist_extended_history: bool, + ) -> CodexResult { + self.spawn_thread_with_source( + config, + initial_history, + Arc::clone(&self.auth_manager), + agent_control, + session_source, + Vec::new(), + persist_extended_history, + None, + ) + .await + } + /// Spawn a new thread with optional history and register it with the manager. #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_thread( diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index d8b83a866..16ff943ab 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents.rs @@ -91,6 +91,7 @@ impl ToolHandler for MultiAgentHandler { mod spawn { use super::*; + use crate::agent::control::SpawnAgentOptions; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; @@ -103,6 +104,8 @@ mod spawn { message: Option, items: Option>, agent_type: Option, + #[serde(default)] + fork_context: bool, } #[derive(Debug, Serialize)] @@ -155,7 +158,7 @@ mod spawn { let result = session .services .agent_control - .spawn_agent( + .spawn_agent_with_options( config, input_items, Some(thread_spawn_source( @@ -163,6 +166,9 @@ mod spawn { child_depth, role_name, )), + SpawnAgentOptions { + fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), + }, ) .await .map_err(collab_spawn_error); diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 47682cf0e..76730d089 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -633,6 +633,15 @@ fn create_spawn_agent_tool(config: &ToolsConfig) -> ToolSpec { )), }, ), + ( + "fork_context".to_string(), + JsonSchema::Boolean { + description: Some( + "When true, fork the current thread history into the new agent before sending the initial prompt. This must be used when you want the new agent to have exactly the same context as you." + .to_string(), + ), + }, + ), ]); ToolSpec::Function(ResponsesApiTool { diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 422510f32..3dc463c4a 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -20,6 +20,8 @@ use tokio::time::sleep; use wiremock::MockServer; const SPAWN_CALL_ID: &str = "spawn-call-1"; +const FORKED_SPAWN_AGENT_OUTPUT_MESSAGE: &str = "You are the newly spawned agent. The prior conversation history was forked from your parent agent. Treat the next user message as your new task, and use the forked history only as background context."; +const TURN_0_FORK_PROMPT: &str = "seed fork context"; const TURN_1_PROMPT: &str = "spawn a child and continue"; const TURN_2_NO_WAIT_PROMPT: &str = "follow up without wait"; const CHILD_PROMPT: &str = "child: do work"; @@ -194,3 +196,116 @@ async fn subagent_notification_is_included_without_wait() -> Result<()> { Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn spawned_child_receives_forked_parent_context() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let seed_turn = mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, TURN_0_FORK_PROMPT), + sse(vec![ + ev_response_created("resp-seed-1"), + ev_assistant_message("msg-seed-1", "seeded"), + ev_completed("resp-seed-1"), + ]), + ) + .await; + + let spawn_args = serde_json::to_string(&json!({ + "message": CHILD_PROMPT, + "fork_context": true, + }))?; + let spawn_turn = mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, TURN_1_PROMPT), + sse(vec![ + ev_response_created("resp-turn1-1"), + ev_function_call(SPAWN_CALL_ID, "spawn_agent", &spawn_args), + ev_completed("resp-turn1-1"), + ]), + ) + .await; + + let _child_request_log = mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, CHILD_PROMPT), + sse(vec![ + ev_response_created("resp-child-1"), + ev_assistant_message("msg-child-1", "child done"), + ev_completed("resp-child-1"), + ]), + ) + .await; + + let _turn1_followup = mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, SPAWN_CALL_ID), + sse(vec![ + ev_response_created("resp-turn1-2"), + ev_assistant_message("msg-turn1-2", "parent done"), + ev_completed("resp-turn1-2"), + ]), + ) + .await; + + let mut builder = test_codex().with_config(|config| { + config.features.enable(Feature::Collab); + }); + let test = builder.build(&server).await?; + + test.submit_turn(TURN_0_FORK_PROMPT).await?; + let _ = seed_turn.single_request(); + + test.submit_turn(TURN_1_PROMPT).await?; + let _ = spawn_turn.single_request(); + + let deadline = Instant::now() + Duration::from_secs(2); + let child_request = loop { + if let Some(request) = server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .find(|request| { + body_contains(request, CHILD_PROMPT) + && body_contains(request, FORKED_SPAWN_AGENT_OUTPUT_MESSAGE) + }) + { + break request; + } + if Instant::now() >= deadline { + anyhow::bail!("timed out waiting for forked child request"); + } + sleep(Duration::from_millis(10)).await; + }; + assert!(body_contains(&child_request, TURN_0_FORK_PROMPT)); + assert!(body_contains(&child_request, "seeded")); + + let child_body = child_request + .body_json::() + .expect("forked child request body should be json"); + let function_call_output = child_body["input"] + .as_array() + .and_then(|items| { + items.iter().find(|item| { + item["type"].as_str() == Some("function_call_output") + && item["call_id"].as_str() == Some(SPAWN_CALL_ID) + }) + }) + .unwrap_or_else(|| panic!("expected forked child request to include spawn_agent output")); + let (content, success) = match &function_call_output["output"] { + serde_json::Value::String(text) => (Some(text.as_str()), None), + serde_json::Value::Object(output) => ( + output.get("content").and_then(serde_json::Value::as_str), + output.get("success").and_then(serde_json::Value::as_bool), + ), + _ => (None, None), + }; + assert_eq!(content, Some(FORKED_SPAWN_AGENT_OUTPUT_MESSAGE)); + assert_ne!(success, Some(false)); + + Ok(()) +}