From 920255e8c69270942206b60f94e26b545af18050 Mon Sep 17 00:00:00 2001 From: Brendan Allan Date: Mon, 16 Feb 2026 12:14:24 +0800 Subject: desktop: use process-wrap instead of manual job object (#13431) --- packages/desktop/src-tauri/src/cli.rs | 162 +++++++++++++++++++++++---- packages/desktop/src-tauri/src/job_object.rs | 145 ------------------------ packages/desktop/src-tauri/src/lib.rs | 15 +-- packages/desktop/src-tauri/src/logging.rs | 11 +- packages/desktop/src-tauri/src/server.rs | 2 +- 5 files changed, 143 insertions(+), 192 deletions(-) delete mode 100644 packages/desktop/src-tauri/src/job_object.rs (limited to 'packages/desktop/src-tauri/src') diff --git a/packages/desktop/src-tauri/src/cli.rs b/packages/desktop/src-tauri/src/cli.rs index dade1a281..0f5cd2ff1 100644 --- a/packages/desktop/src-tauri/src/cli.rs +++ b/packages/desktop/src-tauri/src/cli.rs @@ -1,12 +1,19 @@ use futures::{FutureExt, Stream, StreamExt, future}; +use process_wrap::tokio::CommandWrap; +#[cfg(unix)] +use process_wrap::tokio::ProcessGroup; +#[cfg(windows)] +use process_wrap::tokio::{JobObject, KillOnDrop}; +#[cfg(unix)] +use std::os::unix::process::ExitStatusExt; +use std::{process::Stdio, time::Duration}; use tauri::{AppHandle, Manager, path::BaseDirectory}; -use tauri_plugin_shell::{ - ShellExt, - process::{CommandChild, CommandEvent, TerminatedPayload}, -}; use tauri_plugin_store::StoreExt; use tauri_specta::Event; -use tokio::sync::oneshot; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::Command; +use tokio::sync::{mpsc, oneshot}; +use tokio_stream::wrappers::ReceiverStream; use tracing::Instrument; use crate::constants::{SETTINGS_STORE, WSL_ENABLED_KEY}; @@ -25,6 +32,33 @@ pub struct Config { pub server: Option, } +#[derive(Clone, Debug)] +pub enum CommandEvent { + Stdout(Vec), + Stderr(Vec), + Error(String), + Terminated(TerminatedPayload), +} + +#[derive(Clone, Copy, Debug)] +pub struct TerminatedPayload { + pub code: Option, + pub signal: Option, +} + +#[derive(Clone, Debug)] +pub struct CommandChild { + kill: mpsc::Sender<()>, +} + +impl CommandChild { + pub fn kill(&self) -> std::io::Result<()> { + self.kill + .try_send(()) + .map_err(|e| std::io::Error::other(e.to_string())) + } +} + pub async fn get_config(app: &AppHandle) -> Option { let (events, _) = spawn_command(app, "debug config", &[]).ok()?; @@ -190,7 +224,7 @@ pub fn spawn_command( app: &tauri::AppHandle, args: &str, extra_env: &[(&str, String)], -) -> Result<(impl Stream + 'static, CommandChild), tauri_plugin_shell::Error> { +) -> Result<(impl Stream + 'static, CommandChild), std::io::Error> { let state_dir = app .path() .resolve("", BaseDirectory::AppLocalData) @@ -217,7 +251,7 @@ pub fn spawn_command( .map(|(key, value)| (key.to_string(), value.clone())), ); - let cmd = if cfg!(windows) { + let mut cmd = if cfg!(windows) { if is_wsl_enabled(app) { tracing::info!("WSL is enabled, spawning CLI server in WSL"); let version = app.package_info().version.to_string(); @@ -249,18 +283,16 @@ pub fn spawn_command( script.push(format!("{} exec \"$BIN\" {}", env_prefix.join(" "), args)); - app.shell() - .command("wsl") - .args(["-e", "bash", "-lc", &script.join("\n")]) + let mut cmd = Command::new("wsl"); + cmd.args(["-e", "bash", "-lc", &script.join("\n")]); + cmd } else { - let mut cmd = app - .shell() - .sidecar("opencode-cli") - .unwrap() - .args(args.split_whitespace()); + let sidecar = get_sidecar_path(app); + let mut cmd = Command::new(sidecar); + cmd.args(args.split_whitespace()); for (key, value) in envs { - cmd = cmd.env(key, value); + cmd.env(key, value); } cmd @@ -269,26 +301,111 @@ pub fn spawn_command( let sidecar = get_sidecar_path(app); let shell = get_user_shell(); - let cmd = if shell.ends_with("/nu") { + let line = if shell.ends_with("/nu") { format!("^\"{}\" {}", sidecar.display(), args) } else { format!("\"{}\" {}", sidecar.display(), args) }; - let mut cmd = app.shell().command(&shell).args(["-il", "-c", &cmd]); + let mut cmd = Command::new(shell); + cmd.args(["-il", "-c", &line]); for (key, value) in envs { - cmd = cmd.env(key, value); + cmd.env(key, value); } cmd }; - let (rx, child) = cmd.spawn()?; - let event_stream = tokio_stream::wrappers::ReceiverStream::new(rx); + cmd.stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + + let mut wrap = CommandWrap::from(cmd); + + #[cfg(unix)] + { + wrap.wrap(ProcessGroup::leader()); + } + + #[cfg(windows)] + { + wrap.wrap(JobObject).wrap(KillOnDrop); + } + + let mut child = wrap.spawn()?; + let stdout = child.stdout().take(); + let stderr = child.stderr().take(); + let (tx, rx) = mpsc::channel(256); + let (kill_tx, mut kill_rx) = mpsc::channel(1); + + if let Some(stdout) = stdout { + let tx = tx.clone(); + tokio::spawn(async move { + let mut lines = BufReader::new(stdout).lines(); + while let Ok(Some(line)) = lines.next_line().await { + let _ = tx.send(CommandEvent::Stdout(line.into_bytes())).await; + } + }); + } + + if let Some(stderr) = stderr { + let tx = tx.clone(); + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + while let Ok(Some(line)) = lines.next_line().await { + let _ = tx.send(CommandEvent::Stderr(line.into_bytes())).await; + } + }); + } + + tokio::spawn(async move { + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break Ok(status), + Ok(None) => {} + Err(err) => break Err(err), + } + + tokio::select! { + _ = kill_rx.recv() => { + let _ = child.start_kill(); + } + _ = tokio::time::sleep(Duration::from_millis(100)) => {} + } + }; + + match status { + Ok(status) => { + let payload = TerminatedPayload { + code: status.code(), + signal: signal_from_status(status), + }; + let _ = tx.send(CommandEvent::Terminated(payload)).await; + } + Err(err) => { + let _ = tx.send(CommandEvent::Error(err.to_string())).await; + } + } + }); + + let event_stream = ReceiverStream::new(rx); let event_stream = sqlite_migration::logs_middleware(app.clone(), event_stream); - Ok((event_stream, child)) + Ok((event_stream, CommandChild { kill: kill_tx })) +} + +fn signal_from_status(status: std::process::ExitStatus) -> Option { + #[cfg(unix)] + { + return status.signal(); + } + + #[cfg(not(unix))] + { + let _ = status; + None + } } pub fn serve( @@ -340,7 +457,6 @@ pub fn serve( let _ = tx.send(payload); } } - _ => {} } future::ready(()) diff --git a/packages/desktop/src-tauri/src/job_object.rs b/packages/desktop/src-tauri/src/job_object.rs deleted file mode 100644 index 8d774b14c..000000000 --- a/packages/desktop/src-tauri/src/job_object.rs +++ /dev/null @@ -1,145 +0,0 @@ -//! Windows Job Object for reliable child process cleanup. -//! -//! This module provides a wrapper around Windows Job Objects with the -//! `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` flag set. When the job object handle -//! is closed (including when the parent process exits or crashes), Windows -//! automatically terminates all processes assigned to the job. -//! -//! This is more reliable than manual cleanup because it works even if: -//! - The parent process crashes -//! - The parent is killed via Task Manager -//! - The RunEvent::Exit handler fails to run - -use std::io::{Error, Result}; -#[cfg(windows)] -use std::sync::Mutex; -use windows::Win32::Foundation::{CloseHandle, HANDLE}; -use windows::Win32::System::JobObjects::{ - AssignProcessToJobObject, CreateJobObjectW, JobObjectExtendedLimitInformation, - SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, - JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, -}; -use windows::Win32::System::Threading::{OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE}; - -/// A Windows Job Object configured to kill all assigned processes when closed. -/// -/// When this struct is dropped or when the owning process exits (even abnormally), -/// Windows will automatically terminate all processes that have been assigned to it. -pub struct JobObject(HANDLE); - -// SAFETY: HANDLE is just a pointer-sized value, and Windows job objects -// can be safely accessed from multiple threads. -unsafe impl Send for JobObject {} -unsafe impl Sync for JobObject {} - -impl JobObject { - /// Creates a new anonymous job object with `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` set. - /// - /// When the last handle to this job is closed (including on process exit), - /// Windows will terminate all processes assigned to the job. - pub fn new() -> Result { - unsafe { - // Create an anonymous job object - let job = CreateJobObjectW(None, None).map_err(|e| Error::other(e.message()))?; - - // Configure the job to kill all processes when the handle is closed - let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); - info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; - - SetInformationJobObject( - job, - JobObjectExtendedLimitInformation, - &info as *const _ as *const std::ffi::c_void, - std::mem::size_of::() as u32, - ) - .map_err(|e| Error::other(e.message()))?; - - Ok(Self(job)) - } - } - - /// Assigns a process to this job object by its process ID. - /// - /// Once assigned, the process will be terminated when this job object is dropped - /// or when the owning process exits. - /// - /// # Arguments - /// * `pid` - The process ID of the process to assign - pub fn assign_pid(&self, pid: u32) -> Result<()> { - unsafe { - // Open a handle to the process with the minimum required permissions - // PROCESS_SET_QUOTA and PROCESS_TERMINATE are required by AssignProcessToJobObject - let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, pid) - .map_err(|e| Error::other(e.message()))?; - - // Assign the process to the job - let result = AssignProcessToJobObject(self.0, process); - - // Close our handle to the process - the job object maintains its own reference - let _ = CloseHandle(process); - - result.map_err(|e| Error::other(e.message())) - } - } -} - -impl Drop for JobObject { - fn drop(&mut self) { - unsafe { - // When this handle is closed and it's the last handle to the job, - // Windows will terminate all processes in the job due to KILL_ON_JOB_CLOSE - let _ = CloseHandle(self.0); - } - } -} - -/// Holds the Windows Job Object that ensures child processes are killed when the app exits. -/// On Windows, when the job object handle is closed (including on crash), all assigned -/// processes are automatically terminated by the OS. -#[cfg(windows)] -pub struct JobObjectState { - job: Mutex>, - error: Mutex>, -} - -#[cfg(windows)] -impl JobObjectState { - pub fn new() -> Self { - match JobObject::new() { - Ok(job) => Self { - job: Mutex::new(Some(job)), - error: Mutex::new(None), - }, - Err(e) => { - tracing::error!("Failed to create job object: {e}"); - Self { - job: Mutex::new(None), - error: Mutex::new(Some(format!("Failed to create job object: {e}"))), - } - } - } - } - - pub fn assign_pid(&self, pid: u32) { - if let Some(job) = self.job.lock().unwrap().as_ref() { - if let Err(e) = job.assign_pid(pid) { - tracing::error!(pid, "Failed to assign process to job object: {e}"); - *self.error.lock().unwrap() = - Some(format!("Failed to assign process to job object: {e}")); - } else { - tracing::info!(pid, "Assigned process to job object for automatic cleanup"); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_job_object_creation() { - let job = JobObject::new(); - assert!(job.is_ok(), "Failed to create job object: {:?}", job.err()); - } -} diff --git a/packages/desktop/src-tauri/src/lib.rs b/packages/desktop/src-tauri/src/lib.rs index aa605a923..4a1c8dc4a 100644 --- a/packages/desktop/src-tauri/src/lib.rs +++ b/packages/desktop/src-tauri/src/lib.rs @@ -1,7 +1,5 @@ mod cli; mod constants; -#[cfg(windows)] -mod job_object; #[cfg(target_os = "linux")] pub mod linux_display; mod logging; @@ -10,12 +8,11 @@ mod server; mod window_customizer; mod windows; +use crate::cli::CommandChild; use futures::{ FutureExt, TryFutureExt, future::{self, Shared}, }; -#[cfg(windows)] -use job_object::*; use std::{ env, net::TcpListener, @@ -27,7 +24,6 @@ use std::{ use tauri::{AppHandle, Listener, Manager, RunEvent, State, ipc::Channel}; #[cfg(any(target_os = "linux", all(debug_assertions, windows)))] use tauri_plugin_deep_link::DeepLinkExt; -use tauri_plugin_shell::process::CommandChild; use tauri_specta::Event; use tokio::{ sync::{oneshot, watch}, @@ -631,12 +627,6 @@ async fn initialize(app: AppHandle) { tracing::info!("CLI health check OK"); - #[cfg(windows)] - { - let job_state = app.state::(); - job_state.assign_pid(child.pid()); - } - app.state::().set_child(Some(child)); Ok(ServerReadyData { url, password }) @@ -710,9 +700,6 @@ fn setup_app(app: &tauri::AppHandle, init_rx: watch::Receiver) { #[cfg(any(target_os = "linux", all(debug_assertions, windows)))] app.deep_link().register_all().ok(); - #[cfg(windows)] - app.manage(JobObjectState::new()); - app.manage(InitState { current: init_rx }); } diff --git a/packages/desktop/src-tauri/src/logging.rs b/packages/desktop/src-tauri/src/logging.rs index f794f9c1b..b985b1f9d 100644 --- a/packages/desktop/src-tauri/src/logging.rs +++ b/packages/desktop/src-tauri/src/logging.rs @@ -36,11 +36,7 @@ pub fn init(log_dir: &Path) -> WorkerGuard { tracing_subscriber::registry() .with(filter) .with(fmt::layer().with_writer(std::io::stderr)) - .with( - fmt::layer() - .with_writer(non_blocking) - .with_ansi(false), - ) + .with(fmt::layer().with_writer(non_blocking).with_ansi(false)) .init(); guard @@ -55,10 +51,7 @@ pub fn tail() -> String { return String::new(); }; - let lines: Vec = BufReader::new(file) - .lines() - .map_while(Result::ok) - .collect(); + let lines: Vec = BufReader::new(file).lines().map_while(Result::ok).collect(); let start = lines.len().saturating_sub(TAIL_LINES); lines[start..].join("\n") diff --git a/packages/desktop/src-tauri/src/server.rs b/packages/desktop/src-tauri/src/server.rs index 6dcf0e586..8253482ef 100644 --- a/packages/desktop/src-tauri/src/server.rs +++ b/packages/desktop/src-tauri/src/server.rs @@ -2,12 +2,12 @@ use std::time::{Duration, Instant}; use tauri::AppHandle; use tauri_plugin_dialog::{DialogExt, MessageDialogButtons, MessageDialogResult}; -use tauri_plugin_shell::process::CommandChild; use tauri_plugin_store::StoreExt; use tokio::task::JoinHandle; use crate::{ cli, + cli::CommandChild, constants::{DEFAULT_SERVER_URL_KEY, SETTINGS_STORE, WSL_ENABLED_KEY}, }; -- cgit v1.2.3