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 0424539817..ccd7f20a69 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -46,11 +46,12 @@ async fn handle_spawn_agent( ) -> Result { let ToolInvocation { session, - turn, + step_context, payload, call_id, .. } = invocation; + let turn = &step_context.turn; let arguments = function_arguments(payload)?; let args: SpawnAgentArgs = parse_arguments(&arguments)?; let role_name = args @@ -70,7 +71,7 @@ async fn handle_spawn_agent( } session .emit_turn_item_started( - &turn, + turn, &TurnItem::CollabAgentToolCall(CollabAgentToolCallItem { id: call_id.clone(), tool: CollabAgentTool::SpawnAgent, @@ -128,7 +129,7 @@ async fn handle_spawn_agent( fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), parent_thread_id: Some(session.thread_id), parent_turn_id: Some(turn.sub_id.clone()), - environments: Some(turn.environments.to_selections()), + environments: Some(step_context.environments.to_selections()), }, )) .await @@ -188,7 +189,7 @@ async fn handle_spawn_agent( .unwrap_or_default(); session .emit_turn_item_completed( - &turn, + turn, TurnItem::CollabAgentToolCall(CollabAgentToolCallItem { id: call_id, tool: CollabAgentTool::SpawnAgent, 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 4d1aa3094d..a6b5c46eec 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 @@ -41,12 +41,13 @@ async fn handle_spawn_agent( ) -> Result { let ToolInvocation { session, - turn, + step_context, payload, call_id, source, .. } = invocation; + let turn = &step_context.turn; let arguments = function_arguments(payload)?; let args: SpawnAgentArgs = parse_arguments(&arguments)?; let fork_mode = args.fork_mode()?; @@ -126,7 +127,7 @@ async fn handle_spawn_agent( fork_mode, parent_thread_id: Some(session.thread_id), parent_turn_id: Some(turn.sub_id.clone()), - environments: Some(turn.environments.to_selections()), + environments: Some(step_context.environments.to_selections()), }, ), ) @@ -144,7 +145,7 @@ async fn handle_spawn_agent( .or(spawned_agent.metadata.agent_nickname); emit_sub_agent_activity( &session, - &turn, + turn, SubAgentActivityItem { id: call_id, agent_thread_id: new_thread_id, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 4b8680e56f..6caa24a873 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -62,6 +62,7 @@ use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call; +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; use core_test_support::responses::mount_sse_sequence; @@ -98,6 +99,7 @@ use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; use tempfile::TempDir; +use test_case::test_case; use tokio::net::TcpListener; use tokio::net::TcpStream; use tokio::time::timeout; @@ -787,6 +789,18 @@ async fn serve_environment_with_agents_md( .as_str() .is_some_and(|path| path.ends_with("/AGENTS.md")); let response = match request["method"].as_str() { + Some("environment/info") => json!({ + "id": request["id"], + "result": { "shell": { "name": "zsh", "path": "/bin/zsh" } } + }), + Some("fs/canonicalize") => json!({ + "id": request["id"], + "result": { "path": request["params"]["path"] } + }), + Some("fs/walk") => json!({ + "id": request["id"], + "result": { "entries": [], "errors": [], "truncated": false } + }), Some("fs/getMetadata") if is_agents_md => { json!({ "id": request["id"], @@ -964,6 +978,142 @@ async fn deferred_executor_starts_noise_connection_after_registration() -> Resul Ok(()) } +#[test_case(false, "multi_agent_v1"; "v1")] +#[test_case(true, "collaboration"; "v2")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn deferred_executor_spawn_agent_inherits_ready_step_environments( + multi_agent_v2: bool, + namespace: &str, +) -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let server = start_mock_server().await; + let wait_call_id = "wait-for-spawn-environment"; + let spawn_call_id = "spawn-in-ready-environment"; + let message = "inspect the ready step environment"; + let spawn_arguments = if multi_agent_v2 { + json!({ "message": message, "task_name": "worker" }) + } else { + json!({ "message": message }) + } + .to_string(); + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-wait"), + ev_function_call( + wait_call_id, + "wait_for_environment", + &json!({ "environment_id": REMOTE_ENVIRONMENT_ID }).to_string(), + ), + ev_completed("resp-wait"), + ]), + sse(vec![ + ev_response_created("resp-spawn"), + ev_function_call_with_namespace( + spawn_call_id, + namespace, + "spawn_agent", + &spawn_arguments, + ), + ev_completed("resp-spawn"), + ]), + sse(vec![ + ev_response_created("resp-done-1"), + ev_assistant_message("msg-done-1", "done"), + ev_completed("resp-done-1"), + ]), + sse(vec![ + ev_response_created("resp-done-2"), + ev_assistant_message("msg-done-2", "done"), + ev_completed("resp-done-2"), + ]), + ], + ) + .await; + let mut builder = test_codex() + .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) + .with_config(move |config| { + config.project_doc_max_bytes = 0; + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.enable(Feature::Collab).is_ok()); + if multi_agent_v2 { + assert!(config.features.enable(Feature::MultiAgentV2).is_ok()); + } else { + assert!(config.features.disable(Feature::MultiAgentV2).is_ok()); + } + }); + let (attach_tx, attach_rx) = tokio::sync::oneshot::channel(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let exec_server = tokio::spawn(serve_environment_with_agents_md( + listener, + "", + attach_rx, + shutdown_rx, + )); + let test = timeout( + Duration::from_secs(5), + builder.build_with_remote_and_local_env(&server), + ) + .await + .context("thread startup should not wait for the remote environment")??; + let remote_selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&test.config.cwd), + workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)], + }; + let expected_environments = vec![remote_selection, local(test.config.cwd.clone())]; + let mut created_threads = test.thread_manager.subscribe_thread_created(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "spawn after the environment becomes ready".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + test.config.cwd.clone(), + expected_environments.clone(), + )), + ..Default::default() + }, + }) + .await?; + wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; + attach_tx.send(()).expect("attach remote environment"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + wait_for_response_request_count(&response_mock, /*expected_count*/ 4).await; + + let child_thread_id = timeout(Duration::from_secs(5), created_threads.recv()) + .await + .context("timed out waiting for the subagent thread")??; + let child_thread = test.thread_manager.get_thread(child_thread_id).await?; + assert_eq!( + child_thread.environment_selections().await, + expected_environments + ); + assert!( + response_mock.requests()[1] + .function_call_output_content_and_success(wait_call_id) + .is_some(), + "the spawn request should follow the ready-environment step" + ); + + shutdown_tx + .send(()) + .expect("stop remote environment server"); + exec_server.await?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn deferred_executor_loads_agents_md_when_environment_becomes_ready() -> Result<()> { const AGENTS_CONTENT: &str = "REMOTE_AGENTS_INSTRUCTIONS";