Support remote user shell environment targets

This commit is contained in:
Rasmus Rygaard
2026-06-15 09:54:41 -07:00
parent baddb5e686
commit f854517e24
8 changed files with 395 additions and 16 deletions
@@ -1860,7 +1860,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}")))?;
+13 -4
View File
@@ -290,7 +290,12 @@ pub async fn inter_agent_communication(
}
}
pub async fn run_user_shell_command(sess: &Arc<Session>, sub_id: String, command: String) {
pub async fn run_user_shell_command(
sess: &Arc<Session>,
sub_id: String,
environment_id: Option<String>,
command: String,
) {
if let Some((turn_context, cancellation_token)) =
sess.active_turn_context_and_cancellation_token().await
{
@@ -299,6 +304,7 @@ pub async fn run_user_shell_command(sess: &Arc<Session>, sub_id: String, command
execute_user_shell_command(
session,
turn_context,
environment_id,
command,
cancellation_token,
UserShellCommandMode::ActiveTurnAuxiliary,
@@ -312,7 +318,7 @@ pub async fn run_user_shell_command(sess: &Arc<Session>, sub_id: String, command
sess.spawn_task(
Arc::clone(&turn_context),
Vec::new(),
UserShellCommandTask::new(command),
UserShellCommandTask::new(environment_id, command),
)
.await;
}
@@ -802,8 +808,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 {
+8 -2
View File
@@ -1125,6 +1125,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,
@@ -8334,8 +8335,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();
+336 -2
View File
@@ -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<String>,
command: String,
}
impl UserShellCommandTask {
pub(crate) fn new(command: String) -> Self {
Self { command }
pub(crate) fn new(environment_id: Option<String>, 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<Session>,
turn_context: Arc<TurnContext>,
environment_id: Option<String>,
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.
@@ -338,6 +366,312 @@ pub(crate) async fn execute_user_shell_command(
}
}
async fn execute_remote_user_shell_command(
session: Arc<Session>,
turn_context: Arc<TurnContext>,
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();
let cwd_uri = turn_environment.cwd_uri().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_uri,
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<Shell, String> {
let info = turn_environment
.environment
.info()
.await
.map_err(|err| format!("failed to read environment shell info: {err}"))?;
Shell::from_environment_shell_info(info.shell)
.map_err(|err| format!("unsupported remote shell: {err}"))
}
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<u8>,
stderr: Vec<u8>,
aggregated_output: Vec<u8>,
},
Cancelled,
Error(String),
}
async fn collect_remote_user_shell_output(
exec_backend: Arc<dyn codex_exec_server::ExecBackend>,
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,
+19 -4
View File
@@ -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?;
+3
View File
@@ -649,6 +649,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<String>,
/// The raw command string after '!'
command: String,
},
@@ -1008,7 +1008,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!(
@@ -254,7 +254,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");
@@ -310,7 +313,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");