diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index 0b1f9d10c1..7c94948188 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -40,6 +40,10 @@ impl ToolExecutor for WriteStdinHandler { create_write_stdin_tool() } + fn supports_parallel_tool_calls(&self) -> bool { + true + } + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { Box::pin(self.handle_call(invocation)) } diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 7bedfd43e3..670e94e9b6 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -82,6 +82,7 @@ pub(crate) struct UnifiedExecProcess { output_closed_notify: Arc, cancellation_token: CancellationToken, output_drained: Arc, + interaction_lock: Arc>, state_tx: watch::Sender, state_rx: watch::Receiver, output_task: Option>, @@ -123,6 +124,7 @@ impl UnifiedExecProcess { output_closed_notify, cancellation_token, output_drained, + interaction_lock: Arc::new(Mutex::new(())), state_tx, state_rx, output_task: None, @@ -178,6 +180,10 @@ impl UnifiedExecProcess { Arc::clone(&self.output_drained) } + pub(super) fn interaction_lock(&self) -> Arc> { + Arc::clone(&self.interaction_lock) + } + pub(super) fn has_exited(&self) -> bool { let state = self.state_rx.borrow().clone(); match &self.process_handle { diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index db233397b0..f828ace225 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -659,6 +659,19 @@ impl UnifiedExecProcessManager { ) -> Result { let process_id = request.process_id; + // Different terminal sessions can be polled concurrently, but reads and + // writes against one terminal must not overlap because they share a + // draining output buffer and process lifecycle. + let locked_process = { + let store = self.process_store.lock().await; + let entry = store + .processes + .get(&process_id) + .ok_or(UnifiedExecError::UnknownProcessId { process_id })?; + Arc::clone(&entry.process) + }; + let _interaction_guard = locked_process.interaction_lock().lock_owned().await; + let PreparedProcessHandles { process, output_buffer, @@ -674,7 +687,9 @@ impl UnifiedExecProcessManager { process_id, tty, .. - } = self.prepare_process_handles(process_id).await?; + } = self + .prepare_process_handles(process_id, &locked_process) + .await?; let mut status_after_write = None; if !request.input.is_empty() { @@ -842,12 +857,16 @@ impl UnifiedExecProcessManager { async fn prepare_process_handles( &self, process_id: i32, + expected_process: &Arc, ) -> Result { let mut store = self.process_store.lock().await; let entry = store .processes .get_mut(&process_id) .ok_or(UnifiedExecError::UnknownProcessId { process_id })?; + if !Arc::ptr_eq(&entry.process, expected_process) { + return Err(UnifiedExecError::UnknownProcessId { process_id }); + } entry.last_used = Instant::now(); let OutputHandles { output_buffer, @@ -1355,14 +1374,23 @@ impl UnifiedExecProcessManager { return None; } - let meta: Vec<(i32, Instant, bool)> = store + let mut meta: Vec<(i32, Instant, bool)> = store .processes .iter() .map(|(id, entry)| (*id, entry.last_used, entry.process.has_exited())) .collect(); - if let Some(process_id) = Self::process_id_to_prune_from_meta(&meta) { - return store.remove(process_id); + while let Some(process_id) = Self::process_id_to_prune_from_meta(&meta) { + // Do not prune processes being held by write_stdin. + if let Some(interaction_lock) = store + .processes + .get(&process_id) + .map(|entry| entry.process.interaction_lock()) + && let Ok(_interaction_guard) = interaction_lock.try_lock_owned() + { + return store.remove(process_id); + } + meta.retain(|(id, _, _)| *id != process_id); } None diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 6edc44e297..a20eaa267a 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1109,6 +1109,46 @@ text(JSON.stringify(results)); Ok(()) } +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_write_stdin_calls_run_in_parallel_across_sessions() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn( + &server, + "write to independent terminal sessions in parallel", + r#" +const sessions = await Promise.all([ + tools.exec_command({ cmd: "bash --noprofile --norc", tty: true, yield_time_ms: 250 }), + tools.exec_command({ cmd: "bash --noprofile --norc", tty: true, yield_time_ms: 250 }), +]); + +const results = await Promise.all([ + tools.write_stdin({ + session_id: sessions[0].session_id, + chars: ": > .code-mode-a; while [ ! -e .code-mode-b ]; do sleep 0.01; done; printf 'code-alpha-%s\\n' ready; exit\n", + yield_time_ms: 5000, + }), + tools.write_stdin({ + session_id: sessions[1].session_id, + chars: ": > .code-mode-b; while [ ! -e .code-mode-a ]; do sleep 0.01; done; printf 'code-beta-%s\\n' ready; exit\n", + yield_time_ms: 5000, + }), +]); + +text(JSON.stringify([results[0].output.includes("code-alpha-ready"), results[1].output.includes("code-beta-ready")])); +"#, + ) + .await?; + + let items = custom_tool_output_items(&second_mock.single_request(), "call-1"); + let result: Value = serde_json::from_str(text_item(&items, /*index*/ 1))?; + assert_eq!(result, serde_json::json!([true, true])); + + Ok(()) +} + // This model uses token-based tool-output truncation, giving the downstream // history assertions a stable `…N tokens truncated…` marker. const TOKEN_POLICY_TEST_MODEL: &str = "gpt-5.4"; diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 46b361b73e..62b10cf424 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -3412,6 +3412,78 @@ async fn unified_exec_runs_on_all_platforms() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn write_stdin_calls_run_in_parallel_across_sessions() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_target_windows!(Ok(()), "uses bash and POSIX file rendezvous commands"); + + let server = start_mock_server().await; + let mut builder = test_codex().with_model("gpt-5.4").with_config(|config| { + config.use_experimental_unified_exec_tool = true; + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow feature update"); + }); + let test = builder.build_with_auto_env(&server).await?; + + let start_args = serde_json::to_string(&json!({ + "cmd": "bash --noprofile --norc", + "tty": true, + "yield_time_ms": 250, + }))?; + let cross_session_a = serde_json::to_string(&json!({ + "session_id": 1000, + "chars": ": > .write-stdin-a; while [ ! -e .write-stdin-b ]; do sleep 0.01; done; printf 'alpha-%s\\n' ready; exit\n", + "yield_time_ms": 5_000, + }))?; + let cross_session_b = serde_json::to_string(&json!({ + "session_id": 1001, + "chars": ": > .write-stdin-b; while [ ! -e .write-stdin-a ]; do sleep 0.01; done; printf 'beta-%s\\n' ready; exit\n", + "yield_time_ms": 5_000, + }))?; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-start-1"), + ev_function_call("start-a", "exec_command", &start_args), + ev_function_call("start-b", "exec_command", &start_args), + ev_completed("resp-start-1"), + ]), + sse(vec![ + ev_response_created("resp-cross-1"), + ev_function_call("cross-a", "write_stdin", &cross_session_a), + ev_function_call("cross-b", "write_stdin", &cross_session_b), + ev_completed("resp-cross-1"), + ]), + sse(vec![ + ev_assistant_message("msg-done", "done"), + ev_completed("resp-done"), + ]), + ], + ) + .await; + + submit_unified_exec_turn(&test, "start terminals", PermissionProfile::Disabled).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let outputs = collect_tool_outputs( + &response_mock + .requests() + .into_iter() + .map(|request| request.body_json()) + .collect::>(), + )?; + assert_regex_match(".*alpha-ready.*", &outputs["cross-a"].output); + assert_regex_match(".*beta-ready.*", &outputs["cross-b"].output); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ignore] async fn unified_exec_prunes_exited_sessions_first() -> Result<()> {