diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 72c610ffcc..27996d30a9 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -82,6 +82,17 @@ fn clear_install_in_progress_marker() -> Result<(), String> { clear_install_in_progress_marker_at(&path) } +/// The marker means "a run started and never reported an outcome", so every +/// terminal outcome clears it, failures included. The runtime probe is the +/// backstop for a broken venv, whereas a marker left behind by a failed update +/// pins a working install into a repair it may not be able to finish offline. +/// Never fatal: losing the marker costs a fast path, not correctness. +fn clear_install_marker_best_effort() { + if let Err(msg) = clear_install_in_progress_marker() { + warn!("[install] {}", msg); + } +} + fn clear_install_in_progress_marker_at(path: &Path) -> Result<(), String> { match std::fs::remove_file(path) { Ok(()) => Ok(()), @@ -496,18 +507,7 @@ fn run_install_with_event_mode( ); 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); + warn!("[install] {}", msg); } let (stdout, stderr) = match spawn_script(&script, &args, &state) { @@ -542,20 +542,7 @@ 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); - } + clear_install_marker_best_effort(); diagnostics::finish_attempt( &diagnostics, &attempt, @@ -589,6 +576,7 @@ fn run_install_with_event_mode( let _ = app.emit(event_mode.needs_elevation_event(), &packages); Err("NEEDS_ELEVATION".to_string()) } else { + clear_install_marker_best_effort(); let msg = format!("Installer exited with code {}", code); diagnostics::finish_attempt( &diagnostics, @@ -611,6 +599,7 @@ fn run_install_with_event_mode( Err(msg) } Err(msg) => { + clear_install_marker_best_effort(); diagnostics::finish_attempt(&diagnostics, &attempt, None, false, Some(msg.clone())); clear_current_attempt(&state); if event_mode.emit_terminal_events() { @@ -650,6 +639,9 @@ pub fn record_pending_elevation_canceled( let Some(attempt) = attempt else { return false; }; + // The elevation exit left the marker in place for the resumed run that is + // now not happening. + clear_install_marker_best_effort(); diagnostics::finish_attempt( diagnostics, &attempt, diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 1e129e85b0..e7abaa6715 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -636,6 +636,45 @@ exit 1 } } + #[cfg(unix)] + #[tokio::test] + async fn managed_runtime_probe_survives_more_stdout_than_the_pipe_holds() { + // The probe imports the whole backend, so import-time output can exceed + // the 64 KiB pipe buffer. Waiting on exit before draining wedges the + // child mid-write and reports a healthy install as broken. + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("runtime-noisy"); + remove_managed_capability_cache(); + + let fake = fake_cli( + "runtime-noisy", + r#"#!/bin/sh +if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then + awk 'BEGIN { while (i++ < 4000) print "import chatter on stdout" }' + printf '{"runtime_ready":true}\n' + 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 +fi +exit 1 +"#, + ); + + let started = std::time::Instant::now(); + let probe = probe_managed_bin(fake.bin.clone()).await; + assert!( + matches!(probe, ManagedProbe::Ready { .. }), + "noisy but healthy CLI must probe Ready, got {probe:?}" + ); + assert!( + started.elapsed() < std::time::Duration::from_secs(9), + "probe hit the 10s timeout instead of draining stdout" + ); + remove_managed_capability_cache(); + } + #[cfg(unix)] #[tokio::test] async fn managed_runtime_probe_runs_before_capability_cache() { diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8b53c81..a45131584a 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -306,11 +306,21 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { return Err("studio_runtime_probe_failed".to_string()); }; + // Drain while waiting: the backend import this probe runs can print more than + // the pipe buffer holds, and waiting first deadlocks on it, timing out a + // healthy install into repair. + let reader = tokio::spawn(async move { + let mut buffer = Vec::new(); + let _ = stdout.read_to_end(&mut buffer).await; + buffer + }); + 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; + reader.abort(); info!( "Managed runtime probe timed out in {}ms", started.elapsed().as_millis() @@ -319,10 +329,9 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> { } }; - let mut output = Vec::new(); - if stdout.read_to_end(&mut output).await.is_err() { + let Ok(output) = reader.await else { return Err("studio_runtime_probe_failed".to_string()); - } + }; let payload = String::from_utf8_lossy(&output) .lines() .rev() @@ -389,11 +398,19 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; }; + // Same reason as the runtime probe above. + let reader = tokio::spawn(async move { + let mut buffer = Vec::new(); + let _ = stdout.read_to_end(&mut buffer).await; + buffer + }); + match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) if status.success() => {} Err(_) => { let _ = child.kill().await; let _ = child.wait().await; + reader.abort(); info!( "Managed desktop-capabilities probe timed out in {}ms", started.elapsed().as_millis() @@ -401,6 +418,7 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; } _ => { + reader.abort(); info!( "Managed desktop-capabilities probe exited unsuccessfully in {}ms", started.elapsed().as_millis() @@ -409,10 +427,9 @@ async fn probe_cli_capability(bin: &Path) -> Option { } } - let mut output = Vec::new(); - if stdout.read_to_end(&mut output).await.is_err() { + let Ok(output) = reader.await else { return None; - } + }; let capability = serde_json::from_slice::(&output).ok(); info!( diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index d6ccdf1a83..cf5215e69a 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -272,21 +272,30 @@ def _load_run_module(): def _missing_studio_requirement(run_mod): from importlib.metadata import PackageNotFoundError, distribution - from packaging.requirements import Requirement + from packaging.requirements import InvalidRequirement, 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: + # pip flags (-r, --extra-index-url) are not requirements. Skipping them + # matters: an unparseable line would report the install broken, and + # repair reinstalls the same file, so the failure would never clear. + if not line or line.startswith("-"): + continue + try: + requirement = Requirement(line) + except InvalidRequirement: 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): + # prereleases=True: a prerelease satisfying a floor is not a broken install. + if requirement.specifier and not requirement.specifier.contains( + installed.version, prereleases = True + ): return requirement.name return None diff --git a/unsloth_cli/tests/test_studio_runtime_readiness.py b/unsloth_cli/tests/test_studio_runtime_readiness.py index 23c3ad80b2..9eace51d60 100644 --- a/unsloth_cli/tests/test_studio_runtime_readiness.py +++ b/unsloth_cli/tests/test_studio_runtime_readiness.py @@ -64,6 +64,44 @@ def test_desktop_runtime_check_catches_later_startup_dependency(monkeypatch, cap assert payload["module"] == "definitely-missing-studio-package" +def test_desktop_runtime_check_ignores_pip_flag_lines(monkeypatch, capsys, tmp_path): + """An unparseable line must not fail the check: repair reinstalls the same + file, so the install would be declared broken forever.""" + studio = importlib.import_module("unsloth_cli.commands.studio") + backend = tmp_path / "backend" + requirements = backend / "requirements" + requirements.mkdir(parents = True) + (requirements / "studio.txt").write_text( + "--extra-index-url https://example.invalid/simple\n-r base.txt\n", + encoding = "utf-8", + ) + run_mod = SimpleNamespace(__file__ = str(backend / "run.py")) + monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod) + + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True} + + +def test_desktop_runtime_check_accepts_a_prerelease_over_a_floor(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>=1.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 = "2.0.0b1"), + ) + + studio.desktop_runtime_check(_json_output = True) + + assert json.loads(capsys.readouterr().out) == {"runtime_ready": True} + + def test_desktop_runtime_check_rejects_version_mismatch(monkeypatch, capsys, tmp_path): studio = importlib.import_module("unsloth_cli.commands.studio") backend = tmp_path / "backend"