mirror of
https://github.com/openai/codex.git
synced 2026-09-17 12:23:33 +00:00
R1
This commit is contained in:
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user