From a4964ff46d6c4a8d2e0586a456c2e272483de022 Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Wed, 10 Jun 2026 16:25:10 -0700 Subject: [PATCH] Support remote user shell environment targets --- .../request_processors/thread_processor.rs | 5 +- codex-rs/core/src/session/handlers.rs | 17 +- codex-rs/core/src/session/tests.rs | 10 +- codex-rs/core/src/tasks/user_shell.rs | 348 +++++++++++++++++- codex-rs/core/tests/suite/user_shell_cmd.rs | 23 +- codex-rs/protocol/src/protocol.rs | 3 + .../tui/src/chatwidget/tests/exec_flow.rs | 5 +- .../src/chatwidget/tests/slash_commands.rs | 10 +- 8 files changed, 405 insertions(+), 16 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index fd6432531d..272977c901 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1750,7 +1750,10 @@ impl ThreadRequestProcessor { self.submit_core_op( request_id, thread.as_ref(), - Op::RunUserShellCommand { command }, + Op::RunUserShellCommand { + environment_id: None, + command, + }, ) .await .map_err(|err| internal_error(format!("failed to start shell command: {err}")))?; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index c0f62ec664..344ddb2959 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -333,7 +333,12 @@ pub async fn inter_agent_communication( } } -pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command: String) { +pub async fn run_user_shell_command( + sess: &Arc, + sub_id: String, + environment_id: Option, + command: String, +) { if let Some((turn_context, cancellation_token)) = sess.active_turn_context_and_cancellation_token().await { @@ -342,6 +347,7 @@ pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command execute_user_shell_command( session, turn_context, + environment_id, command, cancellation_token, UserShellCommandMode::ActiveTurnAuxiliary, @@ -355,7 +361,7 @@ pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command sess.spawn_task( Arc::clone(&turn_context), Vec::new(), - UserShellCommandTask::new(command), + UserShellCommandTask::new(environment_id, command), ) .await; } @@ -842,8 +848,11 @@ pub(super) async fn submission_loop( set_thread_memory_mode(&sess, sub.id.clone(), mode).await; false } - Op::RunUserShellCommand { command } => { - run_user_shell_command(&sess, sub.id.clone(), command).await; + Op::RunUserShellCommand { + environment_id, + command, + } => { + run_user_shell_command(&sess, sub.id.clone(), environment_id, command).await; false } Op::ResolveElicitation { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index f6166260a0..a60c257ed4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1127,6 +1127,7 @@ async fn user_shell_commands_do_not_inherit_managed_network_proxy() -> anyhow::R execute_user_shell_command( Arc::clone(&session), turn_context, + None, command, CancellationToken::new(), UserShellCommandMode::StandaloneTurn, @@ -8239,8 +8240,13 @@ async fn run_user_shell_command_does_not_set_reference_context_item() { state.set_reference_context_item(/*item*/ None); } - handlers::run_user_shell_command(&session, "sub-id".to_string(), "echo shell".to_string()) - .await; + handlers::run_user_shell_command( + &session, + "sub-id".to_string(), + None, + "echo shell".to_string(), + ) + .await; let deadline = StdDuration::from_secs(15); let start = std::time::Instant::now(); diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 6affc3d039..5a7084fddb 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -1,10 +1,15 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use codex_async_utils::CancelErr; use codex_async_utils::OrCancelExt; +use codex_exec_server::ExecEnvPolicy as ExecServerEnvPolicy; +use codex_exec_server::ExecOutputStream as ExecServerOutputStream; +use codex_exec_server::ExecParams as ExecServerParams; use codex_network_proxy::PROXY_ACTIVE_ENV_KEY; +use codex_protocol::config_types::ShellEnvironmentPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use tokio_util::sync::CancellationToken; use tracing::error; @@ -29,11 +34,14 @@ use crate::turn_timing::now_unix_timestamp_ms; use crate::user_shell_command::user_shell_command_record_item; use codex_protocol::exec_output::ExecToolCallOutput; use codex_protocol::exec_output::StreamOutput; +use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecCommandBeginEvent; use codex_protocol::protocol::ExecCommandEndEvent; +use codex_protocol::protocol::ExecCommandOutputDeltaEvent; use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::ExecCommandStatus; +use codex_protocol::protocol::ExecOutputStream as ProtocolExecOutputStream; use codex_protocol::protocol::TurnStartedEvent; use codex_sandboxing::SandboxType; use codex_shell_command::parse_command::parse_command; @@ -41,6 +49,7 @@ use codex_shell_command::parse_command::parse_command; use super::SessionTask; use super::SessionTaskContext; use crate::session::session::Session; +use crate::session::turn_context::TurnEnvironment; use codex_protocol::models::PermissionProfile; const USER_SHELL_TIMEOUT_MS: u64 = 60 * 60 * 1000; // 1 hour @@ -57,12 +66,16 @@ pub(crate) enum UserShellCommandMode { #[derive(Clone)] pub(crate) struct UserShellCommandTask { + environment_id: Option, command: String, } impl UserShellCommandTask { - pub(crate) fn new(command: String) -> Self { - Self { command } + pub(crate) fn new(environment_id: Option, command: String) -> Self { + Self { + environment_id, + command, + } } } @@ -85,6 +98,7 @@ impl SessionTask for UserShellCommandTask { execute_user_shell_command( session.clone_session(), turn_context, + self.environment_id.clone(), self.command.clone(), cancellation_token, UserShellCommandMode::StandaloneTurn, @@ -97,6 +111,7 @@ impl SessionTask for UserShellCommandTask { pub(crate) async fn execute_user_shell_command( session: Arc, turn_context: Arc, + environment_id: Option, command: String, cancellation_token: CancellationToken, mode: UserShellCommandMode, @@ -124,6 +139,19 @@ pub(crate) async fn execute_user_shell_command( session.send_event(turn_context.as_ref(), event).await; } + if let Some(environment_id) = environment_id { + execute_remote_user_shell_command( + session, + turn_context, + environment_id, + command, + cancellation_token, + mode, + ) + .await; + return; + } + // Execute the user's script under their default shell when known; this // allows commands that use shell features (pipes, &&, redirects, etc.). // We do not source rc files or otherwise reformat the script. @@ -334,6 +362,322 @@ pub(crate) async fn execute_user_shell_command( } } +async fn execute_remote_user_shell_command( + session: Arc, + turn_context: Arc, + environment_id: String, + command: String, + cancellation_token: CancellationToken, + mode: UserShellCommandMode, +) { + let Some(turn_environment) = turn_context + .environments + .turn_environments + .iter() + .find(|environment| environment.environment_id == environment_id) + .cloned() + else { + let message = format!("unknown turn environment id `{environment_id}`"); + emit_user_shell_failed_without_begin(&session, turn_context.as_ref(), &command, &message) + .await; + return; + }; + + if !turn_environment.environment.is_remote() { + let message = format!("user shell environment `{environment_id}` is not remote"); + emit_user_shell_failed_without_begin(&session, turn_context.as_ref(), &command, &message) + .await; + return; + } + + let remote_shell = match remote_user_shell(&turn_environment).await { + Ok(shell) => shell, + Err(message) => { + emit_user_shell_failed_without_begin( + &session, + turn_context.as_ref(), + &command, + &message, + ) + .await; + return; + } + }; + let use_login_shell = true; + let display_command = remote_shell.derive_exec_args(&command, use_login_shell); + let parsed_cmd = parse_command(&display_command); + let call_id = Uuid::new_v4().to_string(); + let raw_command = command; + let cwd = turn_environment.cwd.clone(); + + session + .send_event( + turn_context.as_ref(), + EventMsg::ExecCommandBegin(ExecCommandBeginEvent { + call_id: call_id.clone(), + process_id: None, + turn_id: turn_context.sub_id.clone(), + started_at_ms: now_unix_timestamp_ms(), + command: display_command.clone(), + cwd: cwd.clone(), + parsed_cmd: parsed_cmd.clone(), + source: ExecCommandSource::UserShell, + interaction_input: None, + }), + ) + .await; + + let process_id = Uuid::new_v4().to_string(); + let exec_params = ExecServerParams { + process_id: codex_exec_server::ProcessId::new(process_id), + argv: display_command.clone(), + cwd: cwd.as_path().to_path_buf(), + env_policy: Some(exec_server_env_policy( + &turn_context.shell_environment_policy, + )), + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + }; + let start = Instant::now(); + let exec_result = collect_remote_user_shell_output( + turn_environment.environment.get_exec_backend(), + exec_params, + StdoutStream { + sub_id: turn_context.sub_id.clone(), + call_id: call_id.clone(), + tx_event: session.get_tx_event(), + }, + cancellation_token, + ) + .await; + + let (exec_output, status) = match exec_result { + RemoteUserShellResult::Output { + exit_code, + stdout, + stderr, + aggregated_output, + } => { + let output = ExecToolCallOutput { + exit_code, + stdout: StreamOutput::new(String::from_utf8_lossy(&stdout).into_owned()), + stderr: StreamOutput::new(String::from_utf8_lossy(&stderr).into_owned()), + aggregated_output: StreamOutput::new( + String::from_utf8_lossy(&aggregated_output).into_owned(), + ), + duration: start.elapsed(), + timed_out: false, + }; + let status = if exit_code == 0 { + ExecCommandStatus::Completed + } else { + ExecCommandStatus::Failed + }; + (output, status) + } + RemoteUserShellResult::Cancelled => { + let aborted_message = "command aborted by user".to_string(); + ( + ExecToolCallOutput { + exit_code: -1, + stdout: StreamOutput::new(String::new()), + stderr: StreamOutput::new(aborted_message.clone()), + aggregated_output: StreamOutput::new(aborted_message), + duration: start.elapsed(), + timed_out: false, + }, + ExecCommandStatus::Failed, + ) + } + RemoteUserShellResult::Error(message) => ( + ExecToolCallOutput { + exit_code: -1, + stdout: StreamOutput::new(String::new()), + stderr: StreamOutput::new(message.clone()), + aggregated_output: StreamOutput::new(message), + duration: start.elapsed(), + timed_out: false, + }, + ExecCommandStatus::Failed, + ), + }; + + session + .send_event( + turn_context.as_ref(), + EventMsg::ExecCommandEnd(ExecCommandEndEvent { + call_id, + process_id: None, + turn_id: turn_context.sub_id.clone(), + completed_at_ms: now_unix_timestamp_ms(), + command: display_command, + cwd, + parsed_cmd, + source: ExecCommandSource::UserShell, + interaction_input: None, + stdout: exec_output.stdout.text.clone(), + stderr: exec_output.stderr.text.clone(), + aggregated_output: exec_output.aggregated_output.text.clone(), + exit_code: exec_output.exit_code, + duration: exec_output.duration, + formatted_output: format_exec_output_str( + &exec_output, + turn_context.truncation_policy, + ), + status, + }), + ) + .await; + persist_user_shell_output( + &session, + turn_context.as_ref(), + &raw_command, + &exec_output, + mode, + ) + .await; +} + +async fn emit_user_shell_failed_without_begin( + session: &Session, + turn_context: &TurnContext, + raw_command: &str, + message: &str, +) { + let exec_output = ExecToolCallOutput { + exit_code: -1, + stdout: StreamOutput::new(String::new()), + stderr: StreamOutput::new(message.to_string()), + aggregated_output: StreamOutput::new(message.to_string()), + duration: Duration::ZERO, + timed_out: false, + }; + persist_user_shell_output( + session, + turn_context, + raw_command, + &exec_output, + UserShellCommandMode::StandaloneTurn, + ) + .await; +} + +async fn remote_user_shell(turn_environment: &TurnEnvironment) -> Result { + let info = turn_environment + .environment + .info() + .await + .map_err(|err| format!("failed to read environment shell info: {err}"))?; + let shell_type = match info.shell.name.as_str() { + "zsh" => crate::shell::ShellType::Zsh, + "bash" => crate::shell::ShellType::Bash, + "powershell" => crate::shell::ShellType::PowerShell, + "sh" => crate::shell::ShellType::Sh, + "cmd" => crate::shell::ShellType::Cmd, + name => return Err(format!("unsupported remote shell `{name}`")), + }; + Ok(Shell { + shell_type, + shell_path: info.shell.path.into(), + shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), + }) +} + +fn exec_server_env_policy(policy: &ShellEnvironmentPolicy) -> ExecServerEnvPolicy { + ExecServerEnvPolicy { + inherit: policy.inherit.clone(), + ignore_default_excludes: policy.ignore_default_excludes, + exclude: policy.exclude.iter().map(ToString::to_string).collect(), + r#set: policy.r#set.clone(), + include_only: policy + .include_only + .iter() + .map(ToString::to_string) + .collect(), + } +} + +enum RemoteUserShellResult { + Output { + exit_code: i32, + stdout: Vec, + stderr: Vec, + aggregated_output: Vec, + }, + Cancelled, + Error(String), +} + +async fn collect_remote_user_shell_output( + exec_backend: Arc, + exec_params: ExecServerParams, + stream: StdoutStream, + cancellation_token: CancellationToken, +) -> RemoteUserShellResult { + let process = match exec_backend.start(exec_params).await { + Ok(started) => started.process, + Err(err) => return RemoteUserShellResult::Error(format!("execution error: {err}")), + }; + let mut after_seq = None; + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let mut aggregated_output = Vec::new(); + let mut exit_code = None; + + loop { + let read = tokio::select! { + _ = cancellation_token.cancelled() => { + let _ = process.terminate().await; + return RemoteUserShellResult::Cancelled; + } + read = process.read(after_seq, None, Some(100)) => read, + }; + let read = match read { + Ok(read) => read, + Err(err) => return RemoteUserShellResult::Error(format!("execution error: {err}")), + }; + after_seq = Some(read.next_seq); + if let Some(code) = read.exit_code { + exit_code = Some(code); + } + if let Some(failure) = read.failure { + return RemoteUserShellResult::Error(failure); + } + for chunk in read.chunks { + let (protocol_stream, target) = match chunk.stream { + ExecServerOutputStream::Stdout | ExecServerOutputStream::Pty => { + (ProtocolExecOutputStream::Stdout, &mut stdout) + } + ExecServerOutputStream::Stderr => (ProtocolExecOutputStream::Stderr, &mut stderr), + }; + let bytes = chunk.chunk.into_inner(); + aggregated_output.extend_from_slice(&bytes); + target.extend_from_slice(&bytes); + let _ = stream + .tx_event + .send(Event { + id: stream.sub_id.clone(), + msg: EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { + call_id: stream.call_id.clone(), + stream: protocol_stream, + chunk: bytes, + }), + }) + .await; + } + if read.closed { + return RemoteUserShellResult::Output { + exit_code: exit_code.unwrap_or(-1), + stdout, + stderr, + aggregated_output, + }; + } + } +} + fn prepare_user_shell_exec_command( display_command: &[String], session_shell: &Shell, diff --git a/codex-rs/core/tests/suite/user_shell_cmd.rs b/codex-rs/core/tests/suite/user_shell_cmd.rs index f7aed6b473..253887dc84 100644 --- a/codex-rs/core/tests/suite/user_shell_cmd.rs +++ b/codex-rs/core/tests/suite/user_shell_cmd.rs @@ -60,7 +60,10 @@ async fn user_shell_cmd_ls_and_cat_in_temp_dir() { // 1) shell command should list the file let list_cmd = "ls".to_string(); codex - .submit(Op::RunUserShellCommand { command: list_cmd }) + .submit(Op::RunUserShellCommand { + environment_id: None, + command: list_cmd, + }) .await .unwrap(); let msg = wait_for_event(&codex, |ev| matches!(ev, EventMsg::ExecCommandEnd(_))).await; @@ -79,7 +82,10 @@ async fn user_shell_cmd_ls_and_cat_in_temp_dir() { // 2) shell command should print the file contents verbatim let cat_cmd = format!("cat {file_name}"); codex - .submit(Op::RunUserShellCommand { command: cat_cmd }) + .submit(Op::RunUserShellCommand { + environment_id: None, + command: cat_cmd, + }) .await .unwrap(); let msg = wait_for_event(&codex, |ev| matches!(ev, EventMsg::ExecCommandEnd(_))).await; @@ -113,7 +119,10 @@ async fn user_shell_cmd_can_be_interrupted() { // Start a long-running command and then interrupt it. let sleep_cmd = "sleep 5".to_string(); codex - .submit(Op::RunUserShellCommand { command: sleep_cmd }) + .submit(Op::RunUserShellCommand { + environment_id: None, + command: sleep_cmd, + }) .await .unwrap(); @@ -216,6 +225,7 @@ async fn user_shell_command_does_not_replace_active_turn() -> anyhow::Result<()> fixture .codex .submit(Op::RunUserShellCommand { + environment_id: None, command: user_shell_command, }) .await?; @@ -281,6 +291,7 @@ async fn user_shell_command_history_is_persisted_and_shared_with_model() -> anyh test.codex .submit(Op::RunUserShellCommand { + environment_id: None, command: command.clone(), }) .await?; @@ -367,7 +378,10 @@ async fn user_shell_command_does_not_set_network_sandbox_env_var() -> anyhow::Re r#"sh -c "printf '%s' \"${CODEX_SANDBOX_NETWORK_DISABLED:-not-set}\"""#.to_string(); test.codex - .submit(Op::RunUserShellCommand { command }) + .submit(Op::RunUserShellCommand { + environment_id: None, + command, + }) .await?; let ExecCommandEndEvent { @@ -409,6 +423,7 @@ async fn user_shell_command_output_is_truncated_in_history() -> anyhow::Result<( test.codex .submit(Op::RunUserShellCommand { + environment_id: None, command: command.clone(), }) .await?; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 5abc9c6399..8b133c258f 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -623,6 +623,9 @@ pub enum Op { /// include shell syntax (pipes, redirects, etc.). Output is streamed via /// `ExecCommand*` events and the UI regains control upon `TurnComplete`. RunUserShellCommand { + /// Optional explicit environment target. `None` preserves the legacy + /// local app-server/TUI path. + environment_id: Option, /// The raw command string after '!' command: String, }, diff --git a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs index 8dbace4378..35dea3d152 100644 --- a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs +++ b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs @@ -992,7 +992,10 @@ async fn bang_shell_enter_while_task_running_submits_run_user_shell_command() { chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); match op_rx.try_recv() { - Ok(Op::RunUserShellCommand { command }) => assert_eq!(command, "echo hi"), + Ok(Op::RunUserShellCommand { + environment_id: None, + command, + }) => assert_eq!(command, "echo hi"), other => panic!("expected RunUserShellCommand op, got {other:?}"), } assert_matches!( diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index 9c2802fb60..2ce99ba7a5 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -232,7 +232,10 @@ async fn queued_bang_shell_dispatches_after_active_turn() { complete_turn_with_message(&mut chat, "turn-1", Some("done")); match op_rx.try_recv() { - Ok(Op::RunUserShellCommand { command }) => assert_eq!(command, "echo hi"), + Ok(Op::RunUserShellCommand { + environment_id: None, + command, + }) => assert_eq!(command, "echo hi"), other => panic!("expected queued shell command op, got {other:?}"), } assert_eq!(next_add_to_history_event(&mut rx), "!echo hi"); @@ -288,7 +291,10 @@ async fn queued_bang_shell_waits_for_user_shell_completion_before_next_input() { complete_turn_with_message(&mut chat, "turn-1", Some("done")); match op_rx.try_recv() { - Ok(Op::RunUserShellCommand { command }) => assert_eq!(command, "echo hi"), + Ok(Op::RunUserShellCommand { + environment_id: None, + command, + }) => assert_eq!(command, "echo hi"), other => panic!("expected queued shell command op, got {other:?}"), } assert_eq!(next_add_to_history_event(&mut rx), "!echo hi");