From a94c5b061b149e88ce5d800ecc3e2ed9126a42db Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Mon, 27 Jul 2026 10:36:22 +0200 Subject: [PATCH] Fix interrupted Studio install recovery --- studio/src-tauri/src/commands.rs | 101 +++-- studio/src-tauri/src/install.rs | 105 +++++ studio/src-tauri/src/preflight.rs | 71 ++- studio/src-tauri/src/preflight/managed.rs | 94 +++- studio/src-tauri/src/process.rs | 417 +++++++++++++----- unsloth_cli/commands/studio.py | 57 +++ .../tests/test_studio_runtime_readiness.py | 85 ++++ 7 files changed, 738 insertions(+), 192 deletions(-) create mode 100644 unsloth_cli/tests/test_studio_runtime_readiness.py diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index 72ea9d6985..b9edcf2da7 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -532,57 +532,70 @@ pub async fn start_managed_repair( let repair_group_id = install::take_pending_repair_group_for_resume(&install_state) .unwrap_or_else(|| diagnostics::begin_repair_group(&diagnostics_state)); - let _ = app.emit("repair-progress", "Updating existing Unsloth install..."); - let update_app = app.clone(); - let update_state = update_state.inner().clone(); - let update_diagnostics = diagnostics_state.clone(); - let update_repair_group_id = repair_group_id.clone(); - let update_result = tokio::task::spawn_blocking(move || { - update::run_backend_update_for_repair( - update_app, - update_state, - update_diagnostics, - update_repair_group_id, - ) - }) - .await - .map_err(|e| format!("Repair update task panicked: {e}"))?; + if install::managed_install_in_progress() { + info!("Interrupted installation detected; skipping incremental repair update"); + let _ = app.emit( + "repair-progress", + "Previous installation was interrupted. Running bundled installer...", + ); + } else { + let _ = app.emit("repair-progress", "Updating existing Unsloth install..."); + let update_app = app.clone(); + let update_state = update_state.inner().clone(); + let update_diagnostics = diagnostics_state.clone(); + let update_repair_group_id = repair_group_id.clone(); + let update_result = tokio::task::spawn_blocking(move || { + update::run_backend_update_for_repair( + update_app, + update_state, + update_diagnostics, + update_repair_group_id, + ) + }) + .await + .map_err(|e| format!("Repair update task panicked: {e}"))?; - match update_result { - Ok(()) if managed_install_ready_after_repair().await => { - info!("Managed repair complete after update"); - diagnostics::finish_repair_group(&diagnostics_state, &repair_group_id, "success", None); - let _ = app.emit("repair-complete", ()); - return Ok(()); - } - Ok(()) => { - warn!("Managed repair update finished, but preflight is still not ready; falling back to installer"); - let _ = app.emit( - "repair-progress", - "Update finished, but Unsloth is still not ready. Running bundled installer...", - ); - } - Err(msg) => { - if msg.to_ascii_lowercase().contains("already running") { - error!("Managed repair update conflict: {}", msg); + match update_result { + Ok(()) if managed_install_ready_after_repair().await => { + info!("Managed repair complete after update"); diagnostics::finish_repair_group( &diagnostics_state, &repair_group_id, - "failed", - Some(msg.clone()), + "success", + None, ); - let _ = app.emit("repair-failed", &msg); - return Err(msg); + let _ = app.emit("repair-complete", ()); + return Ok(()); } + Ok(()) => { + warn!("Managed repair update finished, but preflight is still not ready; falling back to installer"); + let _ = app.emit( + "repair-progress", + "Update finished, but Unsloth is still not ready. Running bundled installer...", + ); + } + Err(msg) => { + if msg.to_ascii_lowercase().contains("already running") { + error!("Managed repair update conflict: {}", msg); + diagnostics::finish_repair_group( + &diagnostics_state, + &repair_group_id, + "failed", + Some(msg.clone()), + ); + let _ = app.emit("repair-failed", &msg); + return Err(msg); + } - warn!( - "Managed repair update failed, falling back to bundled installer: {}", - msg - ); - let _ = app.emit( - "repair-progress", - "Update failed. Running bundled installer...", - ); + warn!( + "Managed repair update failed, falling back to bundled installer: {}", + msg + ); + let _ = app.emit( + "repair-progress", + "Update failed. Running bundled installer...", + ); + } } } diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 024b730735..72c610ffcc 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -7,6 +7,8 @@ use std::process::{Command, ExitStatus, Stdio}; use std::sync::{Arc, Mutex}; use tauri::{AppHandle, Emitter, Manager}; +const INSTALL_IN_PROGRESS_MARKER: &str = ".desktop-install-in-progress"; + // ── Types ── pub struct InstallProcess { @@ -38,6 +40,56 @@ pub fn new_install_state() -> InstallState { use crate::process::trim_line_endings; +fn install_in_progress_marker_path() -> Result { + let home = dirs::home_dir().ok_or("Could not determine home directory")?; + Ok(home + .join(".unsloth") + .join("studio") + .join(INSTALL_IN_PROGRESS_MARKER)) +} + +pub(crate) fn managed_install_in_progress() -> bool { + install_in_progress_marker_path() + .map(|path| path.is_file()) + .unwrap_or(false) +} + +fn create_install_in_progress_marker() -> Result<(), String> { + let path = install_in_progress_marker_path()?; + create_install_in_progress_marker_at(&path) +} + +fn create_install_in_progress_marker_at(path: &Path) -> Result<(), String> { + let parent = path + .parent() + .ok_or_else(|| format!("Invalid install marker path: {}", path.display()))?; + std::fs::create_dir_all(parent) + .map_err(|e| format!("Failed to create {}: {}", parent.display(), e))?; + + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + { + Ok(_) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()), + Err(error) => Err(format!("Failed to create {}: {}", path.display(), error)), + } +} + +fn clear_install_in_progress_marker() -> Result<(), String> { + let path = install_in_progress_marker_path()?; + clear_install_in_progress_marker_at(&path) +} + +fn clear_install_in_progress_marker_at(path: &Path) -> Result<(), String> { + match std::fs::remove_file(path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(format!("Failed to remove {}: {}", path.display(), error)), + } +} + // ── Script Resolution ── /// Returns (script_path, args) depending on dev vs production mode. @@ -443,6 +495,21 @@ fn run_install_with_event_mode( &format!("Using script: {}", script.display()), ); + if let Err(msg) = create_install_in_progress_marker() { + diagnostics::finish_attempt( + &diagnostics, + &attempt, + None, + false, + Some(format!("create_install_marker: {msg}")), + ); + clear_current_attempt(&state); + if event_mode.emit_terminal_events() { + emit_failed(&app, &msg); + } + return Err(msg); + } + let (stdout, stderr) = match spawn_script(&script, &args, &state) { Ok(handles) => handles, Err(msg) => { @@ -475,6 +542,20 @@ fn run_install_with_event_mode( match result { Ok((status, _)) if status.success() => { + if let Err(msg) = clear_install_in_progress_marker() { + diagnostics::finish_attempt( + &diagnostics, + &attempt, + Some(status.to_string()), + false, + Some(format!("clear_install_marker: {msg}")), + ); + clear_current_attempt(&state); + if event_mode.emit_terminal_events() { + emit_failed(&app, &msg); + } + return Err(msg); + } diagnostics::finish_attempt( &diagnostics, &attempt, @@ -877,6 +958,30 @@ fn capped_output_text(bytes: &[u8]) -> String { mod tests { use super::*; + #[test] + fn install_marker_persists_until_explicit_success_cleanup() { + let directory = std::env::temp_dir().join(format!( + "unsloth-install-marker-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let marker = directory.join(INSTALL_IN_PROGRESS_MARKER); + + create_install_in_progress_marker_at(&marker).unwrap(); + assert!(marker.is_file()); + + create_install_in_progress_marker_at(&marker).unwrap(); + assert!(marker.is_file()); + + clear_install_in_progress_marker_at(&marker).unwrap(); + assert!(!marker.exists()); + clear_install_in_progress_marker_at(&marker).unwrap(); + let _ = std::fs::remove_dir_all(directory); + } + #[test] fn elevated_output_cap_is_utf8_boundary_safe() { let text = "é".repeat(40_000); diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 7ef5244754..1e129e85b0 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -566,16 +566,33 @@ mod tests { ( "cap-missing", r#"#!/bin/sh -if [ "$1" = "-h" ]; then exit 0; fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + printf '{"runtime_ready":true}' + exit 0 +fi if [ "$1" = "studio" ] && [ "$2" = "provision-desktop-auth" ] && [ "$3" = "--help" ]; then exit 0; fi exit 1 "#, Some("desktop_capability_probe_failed"), ), + ( + "runtime-missing-dependency", + r#"#!/bin/sh +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + printf '{"runtime_ready":false,"reason":"missing_dependency","module":"structlog"}' + exit 1 +fi +exit 1 +"#, + Some("studio_runtime_missing_dependency"), + ), ( "cap-true-helper-missing", r#"#!/bin/sh -if [ "$1" = "-h" ]; then exit 0; fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + printf '{"runtime_ready":true}' + exit 0 +fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' exit 0 @@ -587,7 +604,10 @@ exit 1 ( "cap-false-helper-ready", r#"#!/bin/sh -if [ "$1" = "-h" ]; then exit 0; fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + printf '{"runtime_ready":true}' + exit 0 +fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":false,"supports_desktop_backend_ownership":true,"desktop_auth_stale_reason":"cap_false","version":"2026.5.3"}' exit 0 @@ -618,26 +638,28 @@ exit 1 #[cfg(unix)] #[tokio::test] - async fn managed_cli_capability_help_probe_runs_before_cache() { + async fn managed_runtime_probe_runs_before_capability_cache() { use std::fs; let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; let _cache_home = ManagedCapabilityCacheHome::new("cache-hit"); remove_managed_capability_cache(); - // `-h` always succeeds unless `modeh` exists; the desktop-capabilities - // probe always succeeds unless `modecap` exists. Toggling those lets us - // prove the ordering: -h runs on every probe (even a cache hit), while - // the heavier capability probe is skipped once the cache is warm. + // Runtime readiness is checked on every probe, while the static + // desktop-capabilities probe is skipped once the cache is warm. let fake = fake_cli( "cap-cache-hit", r#"#!/bin/sh log="$0.calls" -modeh="$0.modeh" +moderuntime="$0.moderuntime" modecap="$0.modecap" printf '%s\n' "$*" >> "$log" -if [ "$1" = "-h" ]; then - if [ -f "$modeh" ]; then exit 42; fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + if [ -f "$moderuntime" ]; then + printf '{"runtime_ready":false,"reason":"missing_dependency","module":"structlog"}' + exit 42 + fi + printf '{"runtime_ready":true}' exit 0 fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then @@ -650,37 +672,42 @@ exit 1 ); let bin = fake.bin.clone(); let calls = bin.with_extension("calls"); - let modeh = bin.with_extension("modeh"); + let moderuntime = bin.with_extension("moderuntime"); let modecap = bin.with_extension("modecap"); - // Cold probe: runs -h and the capability probe, then caches the result. + // Cold probe runs both checks, then caches the static capability. assert!(matches!( probe_managed_bin(bin.clone()).await, ManagedProbe::Ready { .. } )); let first_calls = fs::read_to_string(&calls).unwrap(); - assert!(first_calls.contains("-h")); + assert!(first_calls.contains("studio desktop-runtime-check --json")); assert!(first_calls.contains("studio desktop-capabilities --json")); - // Cache hit: -h still runs, but the capability probe is skipped (breaking - // it via `modecap` proves it is not invoked). + // Cache hit still checks the runtime, but skips the capability probe. fs::write(&modecap, "broken").unwrap(); fs::write(&calls, "").unwrap(); assert!(matches!( probe_managed_bin(bin.clone()).await, ManagedProbe::Ready { .. } )); - assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + assert_eq!( + fs::read_to_string(&calls).unwrap(), + "studio desktop-runtime-check --json\n" + ); - // A non-launchable CLI is caught by the -h probe even with a warm cache: - // preflight reports Stale (for repair) and never trusts the cache. - fs::write(&modeh, "broken").unwrap(); + // A broken runtime is caught even with a warm capability cache. + fs::write(&moderuntime, "broken").unwrap(); fs::write(&calls, "").unwrap(); assert!(matches!( probe_managed_bin(bin).await, - ManagedProbe::Stale { .. } + ManagedProbe::Stale { reason, .. } + if reason == "studio_runtime_missing_dependency" )); - assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + assert_eq!( + fs::read_to_string(&calls).unwrap(), + "studio desktop-runtime-check --json\n" + ); remove_managed_capability_cache(); } diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 0d20f271c5..57b8b53c81 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -36,6 +36,13 @@ struct DesktopCapability { version: Option, } +#[derive(Debug, Deserialize)] +struct DesktopRuntimeCheck { + runtime_ready: bool, + reason: Option, + module: Option, +} + #[derive(Debug, Clone, Deserialize, Serialize)] struct ManagedCapabilityCache { schema: u16, @@ -263,10 +270,12 @@ fn write_cached_capability(fingerprint: &ManagedBinFingerprint, capability: &Des } } -async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { +async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { let started = Instant::now(); let mut cmd = Command::new(bin); - cmd.args(args).stdout(Stdio::null()).stderr(Stdio::null()); + cmd.args(["studio", "desktop-runtime-check", "--json"]) + .stdout(Stdio::piped()) + .stderr(Stdio::null()); #[cfg(target_os = "linux")] if std::env::var_os("APPIMAGE").is_some() { @@ -288,28 +297,60 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { let Ok(mut child) = cmd.spawn() else { info!( - "Managed preflight probe {:?} failed to spawn in {}ms", - args, + "Managed runtime probe failed to spawn in {}ms", started.elapsed().as_millis() ); - return false; + return Err("studio_runtime_probe_failed".to_string()); + }; + let Some(mut stdout) = child.stdout.take() else { + return Err("studio_runtime_probe_failed".to_string()); }; - let ok = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { - Ok(Ok(status)) => status.success(), + let status = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { + Ok(Ok(status)) => status, _ => { let _ = child.kill().await; let _ = child.wait().await; - false + info!( + "Managed runtime probe timed out in {}ms", + started.elapsed().as_millis() + ); + return Err("studio_runtime_probe_failed".to_string()); } }; + + let mut output = Vec::new(); + if stdout.read_to_end(&mut output).await.is_err() { + return Err("studio_runtime_probe_failed".to_string()); + } + let payload = String::from_utf8_lossy(&output) + .lines() + .rev() + .find_map(|line| serde_json::from_str::(line).ok()); + let result = match payload { + Some(payload) if status.success() && payload.runtime_ready => Ok(()), + Some(payload) => { + let reason = match payload.reason.as_deref() { + Some("missing_dependency") => { + info!( + "Managed runtime probe missing dependency module={}", + payload.module.as_deref().unwrap_or("unknown") + ); + "studio_runtime_missing_dependency" + } + Some("backend_import_failed") => "studio_runtime_import_failed", + _ => "studio_runtime_probe_failed", + }; + Err(reason.to_string()) + } + None => Err("studio_runtime_probe_failed".to_string()), + }; info!( - "Managed preflight probe {:?} finished ok={} in {}ms", - args, - ok, + "Managed runtime probe finished ok={} in {}ms", + result.is_ok(), started.elapsed().as_millis() ); - ok + result } async fn probe_cli_capability(bin: &Path) -> Option { @@ -410,22 +451,17 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { let started = Instant::now(); - // Always verify the managed CLI actually launches before trusting the cache. - // A matching capability fingerprint does not prove the binary can still run: - // its venv interpreter or a runtime dependency can be broken while the - // path/size/mtime/markers are unchanged, so the -h probe runs first and a - // non-launchable install is reported Stale for repair. The capability cache - // below still skips the heavier desktop-capabilities probe on a hit. - if !run_cli_probe(&bin, &["-h"]).await { + // Runtime readiness is intentionally uncached. A matching capability + // fingerprint only proves protocol compatibility, not that Studio's backend + // imports are complete after an interrupted dependency transaction. + if let Err(reason) = probe_cli_runtime(&bin).await { info!( - "Managed preflight: cli unusable for {:?} in {}ms", + "Managed preflight: runtime unusable for {:?} reason={} in {}ms", bin, + reason, started.elapsed().as_millis() ); - return ManagedProbe::Stale { - bin, - reason: "cli_unusable".to_string(), - }; + return ManagedProbe::Stale { bin, reason }; } if let Some(fingerprint) = managed_bin_fingerprint(&bin) { @@ -478,6 +514,16 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { pub(super) async fn probe_managed_install() -> ManagedProbe { let started = Instant::now(); let result = match crate::process::find_unsloth_binary() { + Some(bin) if crate::install::managed_install_in_progress() => { + info!( + "Managed preflight: interrupted desktop installation marker found for {:?}", + bin + ); + ManagedProbe::Stale { + bin, + reason: "install_incomplete".to_string(), + } + } Some(bin) => probe_managed_bin(bin).await, None => ManagedProbe::Missing, }; diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index 56d9dd2e21..b7e0a05ee8 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -90,13 +90,18 @@ impl OwnedBackendHandle { } } - fn remove_owner_metadata(self) { + fn cleanup_after_exit(self) { match self { Self::Spawned { - owner: Some(owner), .. + mut child, + owner, + pid, + .. + } => { + terminate_exited_backend_tree(pid, &mut child); + remove_optional_owner(owner); } - | Self::Adopted { owner, .. } => owner.remove(), - Self::Spawned { owner: None, .. } => {} + Self::Adopted { owner, .. } => owner.remove(), } } } @@ -316,6 +321,146 @@ pub(crate) fn trim_line_endings(bytes: &[u8]) -> &[u8] { &bytes[..end] } +#[cfg(unix)] +fn unix_process_group_exists(pid: u32) -> bool { + if pid > i32::MAX as u32 { + return false; + } + let result = unsafe { libc::kill(-(pid as i32), 0) }; + result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM) +} + +#[cfg(unix)] +fn terminate_exited_backend_tree(pid: u32, _child: &mut Box) { + if !unix_process_group_exists(pid) { + return; + } + + info!("Stopping surviving backend process group (pid {})", pid); + unsafe { + libc::kill(-(pid as i32), libc::SIGTERM); + } + for _ in 0..50 { + if !unix_process_group_exists(pid) { + return; + } + std::thread::sleep(Duration::from_millis(100)); + } + + warn!( + "Backend descendants did not exit gracefully, force killing group (pid {})", + pid + ); + unsafe { + libc::kill(-(pid as i32), libc::SIGKILL); + } +} + +#[cfg(windows)] +fn terminate_exited_backend_tree(pid: u32, child: &mut Box) { + info!("Stopping surviving backend job (pid {})", pid); + if let Err(error) = child.kill() { + warn!("Failed to terminate backend job (pid {}): {}", pid, error); + return; + } + if let Err(error) = child.wait() { + warn!("Failed to wait for backend job (pid {}): {}", pid, error); + } +} + +#[cfg(not(any(unix, windows)))] +fn terminate_exited_backend_tree(_pid: u32, _child: &mut Box) {} + +struct BackendExit { + owned: OwnedBackendHandle, + status: String, + intentional: bool, +} + +enum BackendExitPoll { + Running, + Exited(BackendExit), + Stop, + QueryFailed(String), +} + +fn poll_spawned_backend_exit(state: &BackendState, generation: u64) -> BackendExitPoll { + let mut proc = match state.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + }; + if proc.generation != generation { + return BackendExitPoll::Stop; + } + let intentional = proc.intentional_stop; + let status = match proc + .owned + .as_mut() + .and_then(OwnedBackendHandle::spawned_child_mut) + { + Some(child) => match child.try_wait() { + Ok(Some(status)) => status.to_string(), + Ok(None) => return BackendExitPoll::Running, + Err(error) => return BackendExitPoll::QueryFailed(error.to_string()), + }, + None => return BackendExitPoll::Stop, + }; + + let Some(owned) = proc.owned.take() else { + return BackendExitPoll::Stop; + }; + proc.port = None; + proc.diagnostics_session = None; + proc.adopted_watchdog_generation = None; + BackendExitPoll::Exited(BackendExit { + owned, + status, + intentional, + }) +} + +fn monitor_spawned_backend_exit( + app: AppHandle, + state: BackendState, + diagnostics_state: DiagnosticsState, + backend_log: BackendLog, + generation: u64, +) { + let mut query_failure_reported = false; + loop { + match poll_spawned_backend_exit(&state, generation) { + BackendExitPoll::Running => { + query_failure_reported = false; + std::thread::sleep(Duration::from_millis(100)); + } + BackendExitPoll::QueryFailed(error) => { + if !query_failure_reported { + warn!("Failed to query backend process status: {}", error); + query_failure_reported = true; + } + std::thread::sleep(Duration::from_millis(100)); + } + BackendExitPoll::Stop => return, + BackendExitPoll::Exited(exit) => { + info!("Backend exited with status: {}", exit.status); + exit.owned.cleanup_after_exit(); + diagnostics::record_backend_exit( + &diagnostics_state, + &backend_log.session_id, + Some(exit.status), + exit.intentional, + None, + ); + if !exit.intentional { + error!("Backend process exited unexpectedly (crash detected)"); + let _ = app.emit("server-crashed", ()); + } + return; + } + } + } +} + /// Windows `CREATE_NO_WINDOW` flag — suppresses console windows for child processes. #[cfg(windows)] pub(crate) const CREATE_NO_WINDOW: u32 = 0x08000000; @@ -444,6 +589,119 @@ mod tests { ); } + #[cfg(unix)] + #[test] + fn spawned_backend_exit_cleans_surviving_process_group() { + use std::os::unix::process::CommandExt; + + let mut command = Command::new("sh"); + command.args(["-c", "trap '' HUP; sleep 1; sleep 30 & exit 7"]); + command.process_group(0); + let child = command.spawn().unwrap(); + let pid = child.id(); + let state = new_backend_state(); + { + let mut proc = state.lock().unwrap(); + proc.generation = 4; + proc.owned = Some(OwnedBackendHandle::spawned(Box::new(child), None, pid, 4)); + } + + assert!(matches!( + poll_spawned_backend_exit(&state, 4), + BackendExitPoll::Running + )); + + let exit = loop { + match poll_spawned_backend_exit(&state, 4) { + BackendExitPoll::Running => { + std::thread::sleep(Duration::from_millis(50)); + } + BackendExitPoll::Exited(exit) => break exit, + BackendExitPoll::QueryFailed(error) => panic!("{error}"), + BackendExitPoll::Stop => panic!("monitor stopped before observing exit"), + } + }; + assert!(exit.status.contains('7')); + assert!(!exit.intentional); + let group_survived = unix_process_group_exists(pid); + exit.owned.cleanup_after_exit(); + assert!(group_survived); + assert!(!unix_process_group_exists(pid)); + + assert!(matches!( + poll_spawned_backend_exit(&state, 4), + BackendExitPoll::Stop + )); + assert!(!state.lock().unwrap().has_owned_backend()); + } + + #[cfg(windows)] + fn windows_process_is_alive(pid: u32) -> bool { + use windows_sys::Win32::Foundation::{CloseHandle, WAIT_TIMEOUT}; + use windows_sys::Win32::System::Threading::{ + OpenProcess, WaitForSingleObject, PROCESS_SYNCHRONIZE, + }; + + unsafe { + let handle = OpenProcess(PROCESS_SYNCHRONIZE, 0, pid); + if handle.is_null() { + return false; + } + let status = WaitForSingleObject(handle, 0); + let _ = CloseHandle(handle); + status == WAIT_TIMEOUT + } + } + + #[cfg(windows)] + #[test] + fn spawned_backend_exit_cleans_surviving_job() { + let temp = temp_studio_dir("backend-job-cleanup"); + let pid_file = temp.join("descendant.pid"); + let quoted_path = pid_file.to_string_lossy().replace('\'', "''"); + let script = format!( + "$child = Start-Process \"$env:SystemRoot\\System32\\ping.exe\" \ + -ArgumentList '-n','30','127.0.0.1' -WindowStyle Hidden -PassThru; \ + Set-Content -LiteralPath '{quoted_path}' -Value $child.Id -NoNewline; exit 7" + ); + let mut command = Command::new("powershell.exe"); + command.args(["-NoProfile", "-NonInteractive", "-Command", &script]); + let child = spawn_backend_command(command).unwrap(); + let pid = child.id(); + let state = new_backend_state(); + { + let mut proc = state.lock().unwrap(); + proc.generation = 4; + proc.owned = Some(OwnedBackendHandle::spawned(child, None, pid, 4)); + } + + let exit = loop { + match poll_spawned_backend_exit(&state, 4) { + BackendExitPoll::Running => std::thread::sleep(Duration::from_millis(50)), + BackendExitPoll::Exited(exit) => break exit, + BackendExitPoll::QueryFailed(error) => panic!("{error}"), + BackendExitPoll::Stop => panic!("monitor stopped before observing exit"), + } + }; + let descendant_pid = fs::read_to_string(&pid_file) + .unwrap() + .parse::() + .unwrap(); + let descendant_survived = windows_process_is_alive(descendant_pid); + exit.owned.cleanup_after_exit(); + + assert!(exit.status.contains('7')); + assert!(descendant_survived); + for _ in 0..20 { + if !windows_process_is_alive(descendant_pid) { + break; + } + std::thread::sleep(Duration::from_millis(50)); + } + assert!(!windows_process_is_alive(descendant_pid)); + fs::remove_dir_all(temp).unwrap(); + } + fn listening_non_studio_port() -> (u16, mpsc::Sender<()>, std::thread::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); listener.set_nonblocking(true).unwrap(); @@ -548,6 +806,29 @@ fn backend_args(port: u16) -> Vec { .collect() } +fn spawn_backend_command(cmd: Command) -> std::io::Result> { + let mut wrap = CommandWrap::from(cmd); + + #[cfg(windows)] + { + use windows::Win32::System::Threading::{ + CREATE_NEW_PROCESS_GROUP, CREATE_NO_WINDOW as WINDOWS_CREATE_NO_WINDOW, + }; + + // Keep this backend tree independently terminable after its leader exits. + wrap.wrap(CreationFlags( + CREATE_NEW_PROCESS_GROUP | WINDOWS_CREATE_NO_WINDOW, + )); + wrap.wrap(JobObject); + } + + #[cfg(unix)] + wrap.wrap(ProcessGroup::leader()); + + let child: Box = wrap.spawn()?; + Ok(child) +} + /// Spawn the backend process and wire up stdout/stderr reader threads. pub fn start_backend( app: &AppHandle, @@ -634,46 +915,17 @@ pub fn start_backend( let backend_log = diagnostics::begin_backend_session(diagnostics_state, port, generation); - // On Windows, launch the backend directly with hidden-window flags. - // The app process is assigned to a KILL_ON_JOB_CLOSE job in main.rs, so - // children inherit crash-safe cleanup without the buggy per-child JobObject wrapper. - #[cfg(windows)] - let mut child: Box = { - use std::os::windows::process::CommandExt; - - const CREATE_NEW_PROCESS_GROUP: u32 = 0x00000200; - cmd.creation_flags(CREATE_NEW_PROCESS_GROUP | CREATE_NO_WINDOW); - let child = cmd.spawn().map_err(|e| { - let msg = format!("Failed to spawn backend: {}", e); - diagnostics::record_backend_start_failure( - diagnostics_state, - Some(port), - Some(generation), - "spawn_backend", - &msg, - ); - msg - })?; - Box::new(child) - }; - - #[cfg(unix)] - let mut child: Box = { - // Keep the backend tree in a process group on Unix for cleanup. - let mut wrap = CommandWrap::from(cmd); - wrap.wrap(ProcessGroup::leader()); - wrap.spawn().map_err(|e| { - let msg = format!("Failed to spawn backend: {}", e); - diagnostics::record_backend_start_failure( - diagnostics_state, - Some(port), - Some(generation), - "spawn_backend", - &msg, - ); - msg - })? - }; + let mut child = spawn_backend_command(cmd).map_err(|e| { + let msg = format!("Failed to spawn backend: {}", e); + diagnostics::record_backend_start_failure( + diagnostics_state, + Some(port), + Some(generation), + "spawn_backend", + &msg, + ); + msg + })?; let backend_pid = child.id(); let stdout = child.stdout().take(); @@ -731,6 +983,22 @@ pub fn start_backend( }); } + { + let app_handle = app.clone(); + let state_clone = Arc::clone(state); + let diagnostics_clone = diagnostics_state.clone(); + let backend_log_clone = backend_log.clone(); + std::thread::spawn(move || { + monitor_spawned_backend_exit( + app_handle, + state_clone, + diagnostics_clone, + backend_log_clone, + generation, + ); + }); + } + Ok(generation) } @@ -897,7 +1165,7 @@ async fn validate_candidate_port( /// Read lines from a child process stream (stdout or stderr). /// For stdout, parse TAURI_PORT=(\d+) candidates for async validation. -/// When stdout closes and the stop was not intentional, emit server-crashed. +/// Process lifecycle is owned by monitor_spawned_backend_exit. fn read_output_stream( stream: R, app: &AppHandle, @@ -993,63 +1261,6 @@ fn read_output_stream( } } } - - // Stream closed. Only the stdout reader checks for crashes. - if !is_stderr { - let mut exit_record: Option<(String, bool)> = None; - let mut emit_crash = false; - if let Ok(mut proc) = state.lock() { - if proc.generation != generation { - return; - } - let intentional = proc.intentional_stop; - let exited = if let Some(child) = proc - .owned - .as_mut() - .and_then(OwnedBackendHandle::spawned_child_mut) - { - match child.try_wait() { - Ok(Some(status)) => { - info!("Backend stdout stream ended with status: {}", status); - exit_record = Some((status.to_string(), intentional)); - true - } - Ok(None) => { - warn!("Backend stdout stream ended, but process is still running"); - false - } - Err(e) => { - warn!("Failed to query backend status after stdout closed: {}", e); - false - } - } - } else { - false - }; - - if exited { - if let Some(owned) = proc.owned.take() { - owned.remove_owner_metadata(); - } - proc.port = None; - proc.diagnostics_session = None; - emit_crash = !intentional; - } - } - if let Some((status, intentional)) = exit_record { - diagnostics::record_backend_exit( - diagnostics_state, - &backend_log.session_id, - Some(status), - intentional, - None, - ); - } - if emit_crash { - error!("Backend process stdout closed unexpectedly (crash detected)"); - let _ = app.emit("server-crashed", ()); - } - } } fn wait_for_child_exit(child: &mut Box, label: &str) -> bool { @@ -1128,6 +1339,7 @@ fn stop_spawned_backend( && try_exact_port_http_shutdown(port, "Spawned backend") && wait_for_child_exit(&mut child, "Backend") { + terminate_exited_backend_tree(pid, &mut child); remove_optional_owner(owner); return Ok(()); } @@ -1158,6 +1370,7 @@ fn stop_spawned_backend( } if wait_for_child_exit(&mut child, "Backend") { + terminate_exited_backend_tree(pid, &mut child); remove_optional_owner(owner); return Ok(()); } diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index c41963aab9..d6ccdf1a83 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -270,6 +270,27 @@ def _load_run_module(): return _RUN_MODULE +def _missing_studio_requirement(run_mod): + from importlib.metadata import PackageNotFoundError, distribution + from packaging.requirements import Requirement + + requirements = Path(run_mod.__file__).with_name("requirements") / "studio.txt" + for line in requirements.read_text(encoding = "utf-8").splitlines(): + line = line.partition("#")[0].strip() + if not line: + continue + requirement = Requirement(line) + if requirement.marker and not requirement.marker.evaluate(): + continue + try: + installed = distribution(requirement.name) + except PackageNotFoundError: + return requirement.name + if requirement.specifier and not requirement.specifier.contains(installed.version): + return requirement.name + return None + + def _find_setup_script() -> Optional[Path]: """Find studio/setup.sh or studio/setup.ps1. @@ -2809,6 +2830,42 @@ def desktop_capabilities( typer.echo(f"{key}: {value}") +@studio_app.command("desktop-runtime-check", hidden = True) +def desktop_runtime_check( + _json_output: bool = typer.Option( + False, + "--json", + help = "Emit machine-readable JSON.", + ), +): + try: + run_mod = _load_run_module() + missing = _missing_studio_requirement(run_mod) + if missing: + raise ModuleNotFoundError( + f"No distribution named {missing!r}", + name = missing, + ) + except ModuleNotFoundError as exc: + payload = { + "runtime_ready": False, + "reason": "missing_dependency", + "module": exc.name, + } + except Exception as exc: + payload = { + "runtime_ready": False, + "reason": "backend_import_failed", + "error_type": type(exc).__name__, + } + else: + payload = {"runtime_ready": True} + + typer.echo(json.dumps(payload, sort_keys = True)) + if not payload["runtime_ready"]: + raise typer.Exit(code = 1) from None + + @studio_app.command("provision-desktop-auth", hidden = True) def provision_desktop_auth(): """Create/repair desktop auth state for the local machine.""" diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py new file mode 100644 index 0000000000..23c3ad80b2 --- /dev/null +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import importlib +import json +import sys +from types import SimpleNamespace +from pathlib import Path + +import pytest +import typer + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _missing_structlog(): + raise ModuleNotFoundError("No module named 'structlog'", name = "structlog") + + +def test_desktop_runtime_check_reports_missing_dependency_as_json(monkeypatch, capsys): + studio = importlib.import_module("unsloth_cli.commands.studio") + monkeypatch.setattr(studio, "_load_run_module", _missing_structlog) + + with pytest.raises(typer.Exit) as exited: + studio.desktop_runtime_check(_json_output = True) + + assert exited.value.exit_code == 1 + assert json.loads(capsys.readouterr().out) == { + "runtime_ready": False, + "reason": "missing_dependency", + "module": "structlog", + } + + +def test_desktop_runtime_check_reports_success(monkeypatch, capsys): + studio = importlib.import_module("unsloth_cli.commands.studio") + monkeypatch.setattr(studio, "_load_run_module", lambda: object()) + monkeypatch.setattr(studio, "_missing_studio_requirement", lambda _run_mod: None) + + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True} + + +def test_desktop_runtime_check_catches_later_startup_dependency(monkeypatch, capsys, tmp_path): + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text( + "definitely-missing-studio-package\n", + encoding = "utf-8", + ) + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + + with pytest.raises(typer.Exit): + studio.desktop_runtime_check(_json_output = True) + + payload = json.loads(capsys.readouterr().out) + assert payload["module"] == "definitely-missing-studio-package" + + +def test_desktop_runtime_check_rejects_version_mismatch(monkeypatch, capsys, tmp_path): + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text("example-package==2.0\n", encoding = "utf-8") + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + monkeypatch.setattr( + importlib.import_module("importlib.metadata"), + "distribution", + lambda _name: SimpleNamespace(version = "1.0"), + ) + + with pytest.raises(typer.Exit): + studio.desktop_runtime_check(_json_output = True) + + payload = json.loads(capsys.readouterr().out) + assert payload["module"] == "example-package"