This commit is contained in:
jif-oai
2025-10-16 12:28:02 +01:00
parent 65be622e9f
commit d09165383f
10 changed files with 163 additions and 240 deletions

View File

@@ -190,6 +190,7 @@ async fn spawn_with_launch(
program,
args,
env: launch_env,
..
} = launch;
env.extend(launch_env);
@@ -233,7 +234,10 @@ pub(crate) mod errors {
/// error, but the command itself might fail or succeed for other reasons.
/// For now, we conservatively check for well known command failure exit codes and
/// also look for common sandbox denial keywords in the command output.
fn is_likely_sandbox_denied(sandbox_type: SandboxType, exec_output: &ExecToolCallOutput) -> bool {
pub(crate) fn is_likely_sandbox_denied(
sandbox_type: SandboxType,
exec_output: &ExecToolCallOutput,
) -> bool {
if sandbox_type == SandboxType::None || exec_output.exit_code == 0 {
return false;
}

View File

@@ -27,9 +27,12 @@ pub(crate) struct ExecCommandSession {
/// Tracks whether the underlying process has exited.
exit_status: std::sync::Arc<std::sync::atomic::AtomicBool>,
/// Captures the process exit code once it becomes available.
exit_code: std::sync::Arc<StdMutex<Option<i32>>>,
}
impl ExecCommandSession {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
writer_tx: mpsc::Sender<Vec<u8>>,
output_tx: broadcast::Sender<Vec<u8>>,
@@ -38,6 +41,7 @@ impl ExecCommandSession {
writer_handle: JoinHandle<()>,
wait_handle: JoinHandle<()>,
exit_status: std::sync::Arc<std::sync::atomic::AtomicBool>,
exit_code: std::sync::Arc<StdMutex<Option<i32>>>,
) -> (Self, broadcast::Receiver<Vec<u8>>) {
let initial_output_rx = output_tx.subscribe();
(
@@ -49,6 +53,7 @@ impl ExecCommandSession {
writer_handle: StdMutex::new(Some(writer_handle)),
wait_handle: StdMutex::new(Some(wait_handle)),
exit_status,
exit_code,
},
initial_output_rx,
)
@@ -65,6 +70,10 @@ impl ExecCommandSession {
pub(crate) fn has_exited(&self) -> bool {
self.exit_status.load(std::sync::atomic::Ordering::SeqCst)
}
pub(crate) fn exit_code(&self) -> Option<i32> {
self.exit_code.lock().ok().and_then(|guard| *guard)
}
}
impl Drop for ExecCommandSession {

View File

@@ -1,16 +1,7 @@
use std::collections::HashMap;
use std::io::ErrorKind;
use std::io::Read;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU32;
use portable_pty::CommandBuilder;
use portable_pty::PtySize;
use portable_pty::native_pty_system;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::Duration;
use tokio::time::Instant;
@@ -20,6 +11,7 @@ use crate::exec_command::exec_command_params::ExecCommandParams;
use crate::exec_command::exec_command_params::WriteStdinParams;
use crate::exec_command::exec_command_session::ExecCommandSession;
use crate::exec_command::session_id::SessionId;
use crate::pty::spawn_pty_process;
use crate::truncate::truncate_middle;
#[derive(Debug, Default)]
@@ -242,102 +234,12 @@ async fn create_exec_command_session(
login,
} = params;
// Use the native pty implementation for the system
let pty_system = native_pty_system();
// Create a new pty
let pair = pty_system.openpty(PtySize {
rows: 24,
cols: 80,
pixel_width: 0,
pixel_height: 0,
})?;
// Spawn a shell into the pty
let mut command_builder = CommandBuilder::new(shell);
let shell_mode_opt = if login { "-lc" } else { "-c" };
command_builder.arg(shell_mode_opt);
command_builder.arg(cmd);
let args = vec![shell_mode_opt.to_string(), cmd];
let mut child = pair.slave.spawn_command(command_builder)?;
// Obtain a killer that can signal the process independently of `.wait()`.
let killer = child.clone_killer();
// Channel to forward write requests to the PTY writer.
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
// Broadcast for streaming PTY output to readers: subscribers receive from subscription time.
let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256);
// Reader task: drain PTY and forward chunks to output channel.
let mut reader = pair.master.try_clone_reader()?;
let output_tx_clone = output_tx.clone();
let reader_handle = tokio::task::spawn_blocking(move || {
let mut buf = [0u8; 8192];
loop {
match reader.read(&mut buf) {
Ok(0) => break, // EOF
Ok(n) => {
// Forward to broadcast; best-effort if there are subscribers.
let _ = output_tx_clone.send(buf[..n].to_vec());
}
Err(ref e) if e.kind() == ErrorKind::Interrupted => {
// Retry on EINTR
continue;
}
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
// We're in a blocking thread; back off briefly and retry.
std::thread::sleep(Duration::from_millis(5));
continue;
}
Err(_) => break,
}
}
});
// Writer task: apply stdin writes to the PTY writer.
let writer = pair.master.take_writer()?;
let writer = Arc::new(StdMutex::new(writer));
let writer_handle = tokio::spawn({
let writer = writer.clone();
async move {
while let Some(bytes) = writer_rx.recv().await {
let writer = writer.clone();
// Perform blocking write on a blocking thread.
let _ = tokio::task::spawn_blocking(move || {
if let Ok(mut guard) = writer.lock() {
use std::io::Write;
let _ = guard.write_all(&bytes);
let _ = guard.flush();
}
})
.await;
}
}
});
// Keep the child alive until it exits, then signal exit code.
let (exit_tx, exit_rx) = oneshot::channel::<i32>();
let exit_status = Arc::new(AtomicBool::new(false));
let wait_exit_status = exit_status.clone();
let wait_handle = tokio::task::spawn_blocking(move || {
let code = match child.wait() {
Ok(status) => status.exit_code() as i32,
Err(_) => -1,
};
wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);
let _ = exit_tx.send(code);
});
// Create and store the session with channels.
let (session, initial_output_rx) = ExecCommandSession::new(
writer_tx,
output_tx,
killer,
reader_handle,
writer_handle,
wait_handle,
exit_status,
);
Ok((session, initial_output_rx, exit_rx))
let env = HashMap::new();
let spawned = spawn_pty_process(&shell, &args, &env).await?;
Ok((spawned.session, spawned.output_rx, spawned.exit_rx))
}
#[cfg(test)]

View File

@@ -11,6 +11,7 @@ use crate::function_tool::FunctionCallError;
pub(crate) enum ExecutionMode {
Shell,
InteractiveShell,
ApplyPatch(ApplyPatchExec),
}
@@ -36,7 +37,7 @@ static APPLY_PATCH_BACKEND: ApplyPatchBackend = ApplyPatchBackend;
pub(crate) fn backend_for_mode(mode: &ExecutionMode) -> &'static dyn ExecutionBackend {
match mode {
ExecutionMode::Shell => &SHELL_BACKEND,
ExecutionMode::Shell | ExecutionMode::InteractiveShell => &SHELL_BACKEND,
ExecutionMode::ApplyPatch(_) => &APPLY_PATCH_BACKEND,
}
}
@@ -52,7 +53,7 @@ impl ExecutionBackend for ShellBackend {
_config: &ExecutorConfig,
) -> Result<ExecParams, FunctionCallError> {
match mode {
ExecutionMode::Shell => Ok(params),
ExecutionMode::Shell | ExecutionMode::InteractiveShell => Ok(params),
_ => Err(FunctionCallError::RespondToModel(
"shell backend invoked with non-shell mode".to_string(),
)),
@@ -97,9 +98,11 @@ impl ExecutionBackend for ApplyPatchBackend {
justification: params.justification,
})
}
ExecutionMode::Shell => Err(FunctionCallError::RespondToModel(
"apply_patch backend invoked without patch context".to_string(),
)),
ExecutionMode::Shell | ExecutionMode::InteractiveShell => {
Err(FunctionCallError::RespondToModel(
"apply_patch backend invoked without patch context".to_string(),
))
}
}
}

View File

@@ -49,6 +49,7 @@ pub(crate) mod linkers {
pub mod errors {
use crate::error::CodexErr;
use crate::executor::SandboxLaunchError;
use crate::function_tool::FunctionCallError;
use thiserror::Error;
@@ -65,4 +66,10 @@ pub mod errors {
FunctionCallError::RespondToModel(msg.into()).into()
}
}
impl From<SandboxLaunchError> for ExecError {
fn from(err: SandboxLaunchError) -> Self {
CodexErr::from(err).into()
}
}
}

View File

@@ -27,6 +27,8 @@ use crate::executor::sandbox::select_sandbox;
use crate::function_tool::FunctionCallError;
use crate::protocol::AskForApproval;
use crate::protocol::SandboxPolicy;
use crate::pty::SpawnedPty;
use crate::pty::spawn_pty_process;
use crate::shell;
use crate::tools::context::ExecCommandContext;
@@ -124,6 +126,22 @@ impl ExecutionPlan {
&self.request.approval_command
}
pub(crate) async fn spawn_interactive_session(
&self,
session: &Session,
) -> Result<(SpawnedPty, SandboxType), ExecError> {
self.attempt_with_retry(session, |launch| async move {
let sandbox_type = launch.sandbox_type;
let spawned = spawn_pty_process(&launch.program, &launch.args, &launch.env)
.await
.map_err(|err| {
ExecError::rejection(format!("failed to spawn interactive command: {err}"))
})?;
Ok((spawned, sandbox_type))
})
.await
}
pub(crate) async fn prompt_retry_without_sandbox(
&self,
session: &Session,
@@ -221,7 +239,10 @@ impl Executor {
approval_policy: AskForApproval,
context: &ExecCommandContext,
) -> Result<ExecutionPlan, ExecError> {
if matches!(request.mode, ExecutionMode::Shell) {
if matches!(
request.mode,
ExecutionMode::Shell | ExecutionMode::InteractiveShell
) {
request.params =
maybe_translate_shell_command(request.params, session, request.use_shell_profile);
}

View File

@@ -26,6 +26,7 @@ use thiserror::Error;
#[derive(Debug)]
pub(crate) struct SandboxLaunch {
pub sandbox_type: SandboxType,
pub program: String,
pub args: Vec<String>,
pub env: HashMap<String, String>,
@@ -60,6 +61,7 @@ pub(crate) fn build_launch_for_sandbox(
.split_first()
.ok_or(SandboxLaunchError::MissingCommandLine)?;
Ok(SandboxLaunch {
sandbox_type: SandboxType::None,
program: program.clone(),
args: args.to_vec(),
env,
@@ -70,6 +72,7 @@ pub(crate) fn build_launch_for_sandbox(
let args =
create_seatbelt_command_args(command.to_vec(), sandbox_policy, sandbox_policy_cwd);
Ok(SandboxLaunch {
sandbox_type: SandboxType::MacosSeatbelt,
program: MACOS_PATH_TO_SEATBELT_EXECUTABLE.to_string(),
args,
env,
@@ -84,6 +87,7 @@ pub(crate) fn build_launch_for_sandbox(
sandbox_policy_cwd,
);
Ok(SandboxLaunch {
sandbox_type: SandboxType::LinuxSeccomp,
program: exe.to_string_lossy().to_string(),
args,
env,
@@ -196,7 +200,7 @@ pub async fn select_sandbox(
otel_event_manager: &OtelEventManager,
) -> Result<SandboxDecision, ExecError> {
match &request.mode {
ExecutionMode::Shell => {
ExecutionMode::Shell | ExecutionMode::InteractiveShell => {
select_shell_sandbox(
request,
approval_policy,

View File

@@ -39,6 +39,7 @@ mod mcp_tool_call;
mod message_history;
mod model_provider_info;
pub mod parse_command;
mod pty;
mod truncate;
mod unified_exec;
mod user_instructions;

View File

@@ -13,12 +13,18 @@ pub(crate) enum UnifiedExecError {
MissingCommandLine,
#[error("missing codex-linux-sandbox executable path")]
MissingLinuxSandboxExecutable,
#[error("Command denied by sandbox: {message}")]
SandboxDenied { message: String },
}
impl UnifiedExecError {
pub(crate) fn create_session(message: String) -> Self {
Self::CreateSession { message }
}
pub(crate) fn sandbox_denied(message: String) -> Self {
Self::SandboxDenied { message }
}
}
impl From<SandboxLaunchError> for UnifiedExecError {

View File

@@ -1,18 +1,12 @@
use portable_pty::CommandBuilder;
use portable_pty::PtySize;
use portable_pty::native_pty_system;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::io::ErrorKind;
use std::io::Read;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::Ordering;
use tokio::sync::Mutex;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::sync::oneshot::error::TryRecvError;
use tokio::task::JoinHandle;
use tokio::time::Duration;
use tokio::time::Instant;
@@ -20,11 +14,15 @@ use tokio::time::Instant;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::exec::ExecParams;
use crate::exec::ExecToolCallOutput;
use crate::exec::SandboxType;
use crate::exec::StreamOutput;
use crate::exec::is_likely_sandbox_denied;
use crate::exec_command::ExecCommandSession;
use crate::executor::ExecutionMode;
use crate::executor::ExecutionPlan;
use crate::executor::ExecutionRequest;
use crate::executor::SandboxLaunch;
use crate::pty::SpawnedPty;
use crate::tools::context::ExecCommandContext;
use crate::truncate::truncate_middle;
@@ -71,6 +69,7 @@ struct ManagedUnifiedExecSession {
/// `output_buffer`, allowing clients to poll for fresh data.
output_notify: Arc<Notify>,
output_task: JoinHandle<()>,
sandbox_type: SandboxType,
}
#[derive(Debug, Default)]
@@ -110,6 +109,10 @@ impl OutputBufferState {
self.total_bytes = 0;
drained
}
fn snapshot(&self) -> Vec<Vec<u8>> {
self.chunks.iter().cloned().collect()
}
}
type OutputBuffer = Arc<Mutex<OutputBufferState>>;
@@ -119,6 +122,7 @@ impl ManagedUnifiedExecSession {
fn new(
session: ExecCommandSession,
initial_output_rx: tokio::sync::broadcast::Receiver<Vec<u8>>,
sandbox_type: SandboxType,
) -> Self {
let output_buffer = Arc::new(Mutex::new(OutputBufferState::default()));
let output_notify = Arc::new(Notify::new());
@@ -150,6 +154,7 @@ impl ManagedUnifiedExecSession {
output_buffer,
output_notify,
output_task,
sandbox_type,
}
}
@@ -167,6 +172,79 @@ impl ManagedUnifiedExecSession {
fn has_exited(&self) -> bool {
self.session.has_exited()
}
fn exit_code(&self) -> Option<i32> {
self.session.exit_code()
}
async fn snapshot_output(&self) -> Vec<Vec<u8>> {
let guard = self.output_buffer.lock().await;
guard.snapshot()
}
fn sandbox_type(&self) -> SandboxType {
self.sandbox_type
}
async fn check_for_sandbox_denial(&self) -> Result<(), UnifiedExecError> {
if self.sandbox_type() == SandboxType::None || !self.has_exited() {
return Ok(());
}
tokio::time::sleep(Duration::from_millis(20)).await;
let collected_chunks = self.snapshot_output().await;
let mut aggregated: Vec<u8> = Vec::new();
for chunk in collected_chunks {
aggregated.extend_from_slice(&chunk);
}
let aggregated_text = String::from_utf8_lossy(&aggregated).to_string();
let exit_code = self.exit_code().unwrap_or(-1);
let exec_output = ExecToolCallOutput {
exit_code,
stdout: StreamOutput::new(aggregated_text.clone()),
stderr: StreamOutput::new(String::new()),
aggregated_output: StreamOutput::new(aggregated_text.clone()),
duration: Duration::ZERO,
timed_out: false,
};
if is_likely_sandbox_denied(self.sandbox_type(), &exec_output) {
let (snippet, _) = truncate_middle(&aggregated_text, UNIFIED_EXEC_OUTPUT_MAX_BYTES);
let message = if snippet.is_empty() {
format!("exit code {exit_code}")
} else {
snippet
};
return Err(UnifiedExecError::sandbox_denied(message));
}
Ok(())
}
async fn from_spawned(
spawned: SpawnedPty,
sandbox_type: SandboxType,
) -> Result<Self, UnifiedExecError> {
let SpawnedPty {
session,
output_rx,
mut exit_rx,
} = spawned;
let managed = Self::new(session, output_rx, sandbox_type);
let exit_ready = match exit_rx.try_recv() {
Ok(_) | Err(TryRecvError::Closed) => true,
Err(TryRecvError::Empty) => false,
};
if exit_ready {
managed.check_for_sandbox_denial().await?;
}
Ok(managed)
}
}
impl Drop for ManagedUnifiedExecSession {
@@ -180,13 +258,7 @@ impl UnifiedExecSessionManager {
&self,
command: Vec<String>,
context: &UnifiedExecContext<'_>,
) -> Result<
(
ExecCommandSession,
tokio::sync::broadcast::Receiver<Vec<u8>>,
),
UnifiedExecError,
> {
) -> Result<ManagedUnifiedExecSession, UnifiedExecError> {
let executor = &context.session.services.executor;
let otel_event_manager = context.turn.client.get_otel_event_manager();
let approval_command = command.clone();
@@ -210,7 +282,7 @@ impl UnifiedExecSessionManager {
justification: None,
},
approval_command,
mode: ExecutionMode::Shell,
mode: ExecutionMode::InteractiveShell,
stdout_stream: None,
use_shell_profile: false,
};
@@ -231,10 +303,12 @@ impl UnifiedExecSessionManager {
.await
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
plan.attempt_with_retry(context.session, |launch| async move {
create_unified_exec_session(&launch).await
})
.await
let (spawned, sandbox_type) = plan
.spawn_interactive_session(context.session)
.await
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
ManagedUnifiedExecSession::from_spawned(spawned, sandbox_type).await
}
pub async fn handle_request(
@@ -285,9 +359,7 @@ impl UnifiedExecSessionManager {
} else {
let command = request.input_chunks.to_vec();
let new_id = self.next_session_id.fetch_add(1, Ordering::SeqCst);
let (session, initial_output_rx) =
self.open_session_with_sandbox(command, &context).await?;
let managed_session = ManagedUnifiedExecSession::new(session, initial_output_rx);
let managed_session = self.open_session_with_sandbox(command, &context).await?;
let (buffer, notify) = managed_session.output_handles();
writer_tx = managed_session.writer_sender();
output_buffer = buffer;
@@ -388,112 +460,6 @@ impl UnifiedExecSessionManager {
}
}
async fn create_unified_exec_session(
launch: &SandboxLaunch,
) -> Result<
(
ExecCommandSession,
tokio::sync::broadcast::Receiver<Vec<u8>>,
),
UnifiedExecError,
> {
if launch.program.is_empty() {
return Err(UnifiedExecError::MissingCommandLine);
}
let pty_system = native_pty_system();
let pair = pty_system
.openpty(PtySize {
rows: 24,
cols: 80,
pixel_width: 0,
pixel_height: 0,
})
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
// Safe thanks to the check at the top of the function.
let mut command_builder = CommandBuilder::new(launch.program.clone());
for arg in &launch.args {
command_builder.arg(arg.clone());
}
for (key, value) in &launch.env {
command_builder.env(key.clone(), value.clone());
}
let mut child = pair
.slave
.spawn_command(command_builder)
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
let killer = child.clone_killer();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
let (output_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(256);
let mut reader = pair
.master
.try_clone_reader()
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
let output_tx_clone = output_tx.clone();
let reader_handle = tokio::task::spawn_blocking(move || {
let mut buf = [0u8; 8192];
loop {
match reader.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
let _ = output_tx_clone.send(buf[..n].to_vec());
}
Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
std::thread::sleep(Duration::from_millis(5));
continue;
}
Err(_) => break,
}
}
});
let writer = pair
.master
.take_writer()
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
let writer = Arc::new(StdMutex::new(writer));
let writer_handle = tokio::spawn({
let writer = writer.clone();
async move {
while let Some(bytes) = writer_rx.recv().await {
let writer = writer.clone();
let _ = tokio::task::spawn_blocking(move || {
if let Ok(mut guard) = writer.lock() {
use std::io::Write;
let _ = guard.write_all(&bytes);
let _ = guard.flush();
}
})
.await;
}
}
});
let exit_status = Arc::new(AtomicBool::new(false));
let wait_exit_status = Arc::clone(&exit_status);
let wait_handle = tokio::task::spawn_blocking(move || {
let _ = child.wait();
wait_exit_status.store(true, Ordering::SeqCst);
});
let (session, initial_output_rx) = ExecCommandSession::new(
writer_tx,
output_tx,
killer,
reader_handle,
writer_handle,
wait_handle,
exit_status,
);
Ok((session, initial_output_rx))
}
#[cfg(test)]
#[cfg(unix)]
mod tests {