diff --git a/studio/backend/main.py b/studio/backend/main.py index 731625ca74..0a5b775775 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -441,9 +441,30 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None: ).start() +def _warm_rag_embedder() -> None: + """Warm RAG embeddings without blocking backend readiness.""" + try: + from storage import rag_db + + if not rag_db.RAG_AVAILABLE: + return + from core.rag import embeddings + + embeddings.warm() + except Exception: + pass + + @asynccontextmanager async def lifespan(app: FastAPI): """Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache.""" + + import time as _time + + _lifespan_started = _time.perf_counter() + import structlog as _structlog + + _lifespan_log = _structlog.get_logger(__name__) clear_unsloth_compiled_cache() # Remove stale .venv_overlay from old versions; switching now uses .venv_t5/. @@ -454,6 +475,11 @@ async def lifespan(app: FastAPI): # Detect hardware first — sets the DEVICE global used everywhere. detect_hardware() + _lifespan_log.info( + "lifespan hardware detection completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) + # Apple Silicon with MLX missing => Train/Export are greyed out (chat-only). # Reinstall mlx by name on a background thread (off the critical path) and # re-detect, so a reinstall/update that dropped mlx self-heals. No-op @@ -465,7 +491,13 @@ async def lifespan(app: FastAPI): import structlog as _structlog _structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc) - # Reap download workers orphaned by a previous crash before new downloads start. + # Reap workers/runs orphaned by a previous crash before new work starts. + try: + from storage.studio_db import cleanup_orphaned_runs + cleanup_orphaned_runs() + except Exception as exc: + _lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc) + reap_hub_orphan_workers() # llama.cpp probes: capability (MTP support) + freshness (release age). @@ -479,45 +511,23 @@ async def lifespan(app: FastAPI): app.state.llama_cpp_freshness = None _start_llama_cpp_probes_if_enabled(app) - from storage.studio_db import cleanup_orphaned_runs - - try: - cleanup_orphaned_runs() - except Exception as exc: - import structlog - structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) - - # Same for RAG: fail ingestion jobs stranded mid-ingest by a crash. try: from storage.rag_db import reconcile_orphaned_ingestion_jobs reconcile_orphaned_ingestion_jobs() except Exception as exc: - import structlog - structlog.get_logger(__name__).warning( - "reconcile_orphaned_ingestion_jobs failed at startup: %s", exc - ) + _lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc) _start_helper_precache_if_enabled() + threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() - # Warm the RAG embedder so the first upload skips the cold load. Non-fatal. - def _warm_rag_embedder(): - try: - from storage import rag_db - - if not rag_db.RAG_AVAILABLE: - return - from core.rag import embeddings - - embeddings.warm() - except Exception: - pass - - threading.Thread(target = _warm_rag_embedder, daemon = True).start() - - # Initialize RSA key pair for API key encryption (external providers) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair init_key_pair() + _lifespan_log.info( + "lifespan pre-auth setup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() @@ -532,6 +542,11 @@ async def lifespan(app: FastAPI): print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() + + _lifespan_log.info( + "lifespan startup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) yield from core.inference.llama_http import aclose as _close_llama_http @@ -919,6 +934,21 @@ install_api_error_handlers(app) # ============ Health and System Endpoints ============ +@app.get("/api/liveness") +async def liveness_check(): + """Cheap process liveness for desktop port validation.""" + return { + "status": "alive", + "service": "Unsloth UI Backend", + "desktop_protocol_version": 1, + "desktop_manageability_version": 1, + "supports_desktop_auth": True, + "supports_desktop_backend_ownership": True, + "studio_root_id": _studio_root_id(), + **({"desktop_owner": owner} if (owner := _desktop_owner()) else {}), + } + + @app.get("/api/health") async def health_check(request: Request): """Liveness plus launcher capability bits; host fingerprint gated on a bearer. diff --git a/studio/backend/run.py b/studio/backend/run.py index d4cbc26b41..6f36f0928a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -933,6 +933,9 @@ def run_server( """ global _server, _server_thread, _shutdown_event + boot_started = time.perf_counter() + logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port) + # Reap every child if the parent dies abnormally (terminal close, Task # Manager kill, SIGKILL); must run before any child can spawn. from utils.process_lifetime import initialize_parent_lifetime @@ -984,7 +987,14 @@ def run_server( from threading import Thread, Event import uvicorn + import_started = time.perf_counter() + from main import app, setup_frontend, _IS_COLAB + + logger.info( + "Imported FastAPI app in %.1fms", + (time.perf_counter() - import_started) * 1000, + ) from utils.paths import ensure_studio_directories # Allow local stdio MCP servers on a loopback bind (the user's own machine), @@ -997,6 +1007,11 @@ def run_server( # Create all standard directories on startup. ensure_studio_directories() + logger.info( + "Ensured Studio directories in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + # Auto-find a free port if the requested one is in use. if not _is_port_free(host, port): original_port = port @@ -1060,6 +1075,11 @@ def run_server( display_host = _resolve_external_ip() if host == "0.0.0.0" else host _install_uvicorn_startup_log_rewrite(host, display_host) + logger.info( + "run_server pre-uvicorn setup completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + ready_event = Event() startup_failed = Event() startup_errors = [] @@ -1068,6 +1088,10 @@ def run_server( async def startup(self, *args, **kwargs): await super().startup(*args, **kwargs) if getattr(self, "started", False) and not self.should_exit: + logger.info( + "Uvicorn startup hook completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) ready_event.set() # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. @@ -1150,6 +1174,11 @@ def run_server( _shutdown_event.set() raise + logger.info( + "run_server uvicorn ready after %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + _write_pid_file() import atexit diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 914abbbf1d..176665769d 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -354,12 +354,11 @@ function TauriWrapper({ children }: { children: ReactNode }) { ); } - const showApp = status === "running" && desktopAuthReady; + const showApp = status === "running"; + const desktopBooting = status === "running" && !desktopAuthReady; + const showInteractiveApp = showApp && desktopAuthReady; const startupStatus = status === "running" ? "starting" : status; - const startupProgressDetail = - status === "running" && !desktopAuthReady - ? "Signing in to desktop session..." - : progressDetail; + const startupProgressDetail = progressDetail; const usesCustomTitlebar = shouldUseCustomWindowTitlebar(); const usesNativeMacTitlebar = shouldUseNativeMacWindowTitlebar(); const hidesTitlebarSidebar = HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); @@ -369,12 +368,23 @@ function TauriWrapper({ children }: { children: ReactNode }) { - + {showInteractiveApp ? : null} - - {children} + {showInteractiveApp ? : null} + {showInteractiveApp ? children : null} + {desktopBooting ? ( +
+
+
Preparing Studio
+
The local backend is ready. Signing in to your desktop session before loading chats.
+
+
+ Signing in to desktop session... +
+
+ ) : null} ) : ( number | null, shouldContinue: () => boolean, ): Promise { @@ -91,15 +90,7 @@ async function waitForManagedServerReady( continue; } - const healthy = await invoke("check_health", { port }); - if (!shouldContinue()) { - return { status: "aborted" }; - } - if (healthy && getPort() === port) { - return { status: "ready", port }; - } - - await wait(MANAGED_STARTUP_POLL_MS); + return { status: "ready", port }; } } @@ -280,10 +271,9 @@ export function useTauriBackend() { // backend/run.py keeps the 8888-8908 fallback via server-port/TAURI_PORT. await invoke("start_managed_server", { port: 8888 }); - // Wait for the owned backend's server-port event. Don't attach to an - // external backend if the managed start doesn't report a port. - const startupResult = await waitForManagedServerReady( - invoke, + // Rust emits server-port only after validating the desktop-owned process. + // Treat that as the UI handoff point instead of doing a second health poll. + const startupResult = await waitForManagedServerPort( () => portRef.current, () => startingRef.current, ); diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index 7c1ce611b8..6bc2116786 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -65,10 +65,18 @@ pub async fn desktop_preflight( shutdown: tauri::State<'_, ShutdownFlag>, diagnostics: tauri::State<'_, DiagnosticsState>, ) -> Result { + let started = Instant::now(); let (result, adopted_watchdog_generation) = crate::preflight::desktop_preflight_result_with_state(state.inner()).await?; diagnostics::record_preflight(&diagnostics, &result); + info!( + "desktop_preflight completed disposition={:?} port={:?} in {}ms", + result.disposition, + result.port, + started.elapsed().as_millis() + ); + if let Some((generation, newly_adopted)) = adopted_watchdog_generation { if newly_adopted { if let Some(port) = result.port { @@ -205,9 +213,17 @@ pub async fn start_managed_server( port: u16, ) -> Result<(), String> { info!("start_managed_server command called with port {}", port); + + let started = Instant::now(); let diagnostics_state = diagnostics.inner().clone(); let generation = process::start_backend(&app, &state, port, &shutdown, &diagnostics_state)?; + info!( + "start_managed_server spawned generation={} in {}ms", + generation, + started.elapsed().as_millis() + ); + let watchdog_state = state.inner().clone(); let watchdog_shutdown = shutdown.inner().clone(); let watchdog_app = app.clone(); diff --git a/studio/src-tauri/src/desktop_backend_owner.rs b/studio/src-tauri/src/desktop_backend_owner.rs index ed4b5f7fd9..c7d0a7b309 100644 --- a/studio/src-tauri/src/desktop_backend_owner.rs +++ b/studio/src-tauri/src/desktop_backend_owner.rs @@ -93,7 +93,7 @@ enum PreviousAppPidStatus { Uncertain, } -#[derive(Debug, Deserialize)] +#[derive(Clone, Debug, Deserialize)] struct HealthDesktopOwner { kind: Option, token_sha256: Option, @@ -101,15 +101,7 @@ struct HealthDesktopOwner { #[derive(Debug, Deserialize)] struct HealthResponse { - status: Option, - service: Option, version: Option, - desktop_protocol_version: Option, - desktop_manageability_version: Option, - supports_desktop_auth: Option, - supports_desktop_backend_ownership: Option, - studio_root_id: Option, - desktop_owner: Option, } #[derive(Debug)] @@ -123,6 +115,18 @@ struct DesktopLoginPayload<'a> { secret: &'a str, } +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DesktopLiveness { + status: Option, + service: Option, + desktop_protocol_version: Option, + desktop_manageability_version: Option, + supports_desktop_auth: Option, + supports_desktop_backend_ownership: Option, + studio_root_id: Option, + desktop_owner: Option, +} + #[derive(Deserialize)] struct TokenResponse { access_token: String, @@ -290,10 +294,10 @@ impl BackendOwnerState { } pub(crate) fn verifies_exact_port_blocking(&self, port: u16) -> bool { - match fetch_health_blocking(port) { - Ok(Some(health)) => { - health_verifies_metadata(&health, &self.metadata) - && lifecycle_control_block_reason(&health).is_none() + match fetch_liveness_blocking(port) { + Ok(Some(liveness)) => { + liveness_verifies_metadata(&liveness, &self.metadata) + && lifecycle_control_block_reason(&liveness).is_none() } _ => false, } @@ -498,66 +502,110 @@ pub(crate) fn test_owner_state(root_id: &str, token: &str, port: u16) -> Backend } } -fn health_verifies_metadata(health: &HealthResponse, metadata: &DesktopBackendMetadata) -> bool { - let healthy = health.status.as_deref() == Some("healthy") - && health.service.as_deref() == Some("Unsloth UI Backend"); - let Some(owner) = health.desktop_owner.as_ref() else { +fn liveness_verifies_metadata( + liveness: &DesktopLiveness, + metadata: &DesktopBackendMetadata, +) -> bool { + let alive = matches!(liveness.status.as_deref(), Some("alive") | Some("healthy")) + && liveness.service.as_deref() == Some("Unsloth UI Backend"); + let Some(owner) = liveness.desktop_owner.as_ref() else { return false; }; - healthy + alive && owner_matches_metadata( metadata, - health.studio_root_id.as_deref(), + liveness.studio_root_id.as_deref(), owner.kind.as_deref(), owner.token_sha256.as_deref(), ) } -fn lifecycle_control_block_reason(health: &HealthResponse) -> Option { - if health.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) { +fn lifecycle_control_block_reason(liveness: &DesktopLiveness) -> Option { + if liveness.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) { return Some("desktop_protocol_incompatible".to_string()); } - if health.supports_desktop_auth != Some(true) { + if liveness.supports_desktop_auth != Some(true) { return Some("desktop_auth_unsupported".to_string()); } - if health.desktop_manageability_version.unwrap_or(0) + if liveness.desktop_manageability_version.unwrap_or(0) < crate::preflight::DESKTOP_MANAGEABILITY_VERSION { return Some("desktop_manageability_unsupported".to_string()); } - if health.supports_desktop_backend_ownership != Some(true) { + if liveness.supports_desktop_backend_ownership != Some(true) { return Some("desktop_backend_ownership_unsupported".to_string()); } None } -fn ready_for_use_status(health: &HealthResponse) -> OwnedBackendReadiness { - match crate::preflight::backend_version_stale_reason(health.version.as_deref()) { +fn ready_for_use_status(health: Option<&HealthResponse>) -> OwnedBackendReadiness { + let version = health + .and_then(|h| h.version.as_deref()) + .filter(|v| !v.is_empty()); + match crate::preflight::backend_version_stale_reason(version) { Some(reason) => OwnedBackendReadiness::Stale { reason }, None => OwnedBackendReadiness::Ready, } } -async fn fetch_health(port: u16) -> Result, reqwest::Error> { +async fn health_ready_status(port: u16) -> OwnedBackendReadiness { + match fetch_health(port).await { + Ok(health) => ready_for_use_status(health.as_ref()), + Err(reason) => OwnedBackendReadiness::Stale { reason }, + } +} + +async fn fetch_liveness(port: u16) -> Result, reqwest::Error> { let client = reqwest::Client::builder() .timeout(LOCAL_HTTP_TIMEOUT) .build()?; + for path in ["/api/liveness", "/api/health"] { + let response = client + .get(format!("http://127.0.0.1:{port}{path}")) + .send() + .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" { + continue; + } + if !response.status().is_success() { + return Ok(None); + } + return response.json::().await.map(Some); + } + Ok(None) +} + +fn fetch_liveness_blocking(port: u16) -> Result, String> { + for path in ["/api/liveness", "/api/health"] { + let response = http_request_blocking(port, "GET", path, &[], &[])?; + if response.status == 404 && path == "/api/liveness" { + continue; + } + if !(200..300).contains(&response.status) { + return Ok(None); + } + return serde_json::from_slice::(&response.body) + .map(Some) + .map_err(|e| e.to_string()); + } + Ok(None) +} +async fn fetch_health(port: u16) -> Result, String> { + let client = reqwest::Client::builder() + .timeout(LOCAL_HTTP_TIMEOUT) + .build() + .map_err(|e| e.to_string())?; let response = client .get(format!("http://127.0.0.1:{port}/api/health")) .send() - .await?; + .await + .map_err(|e| e.to_string())?; if !response.status().is_success() { return Ok(None); } - response.json::().await.map(Some) -} - -fn fetch_health_blocking(port: u16) -> Result, String> { - let response = http_request_blocking(port, "GET", "/api/health", &[], &[])?; - if !(200..300).contains(&response.status) { - return Ok(None); - } - serde_json::from_slice::(&response.body) + response + .json::() + .await .map(Some) .map_err(|e| e.to_string()) } @@ -618,21 +666,21 @@ pub(crate) async fn probe_owned_backend_state( }; let mut verified = Vec::new(); for port in ports { - let health = match fetch_health(port).await { - Ok(Some(health)) => health, + let liveness = match fetch_liveness(port).await { + Ok(Some(liveness)) => liveness, Ok(None) => continue, Err(error) => { warn!( - "Desktop-owned backend probe skipped port {} after health error: {}", + "Desktop-owned backend probe skipped port {} after liveness error: {}", port, error ); continue; } }; - if !health_verifies_metadata(&health, &owner.metadata) { + if !liveness_verifies_metadata(&liveness, &owner.metadata) { continue; } - if let Some(reason) = lifecycle_control_block_reason(&health) { + if let Some(reason) = lifecycle_control_block_reason(&liveness) { return OwnedBackendProbe::Unmanageable { port, reason }; } if !desktop_login_route_compatible(port).await { @@ -646,7 +694,7 @@ pub(crate) async fn probe_owned_backend_state( return OwnedBackendProbe::Unmanageable { port, reason }; } } - verified.push((port, ready_for_use_status(&health))); + verified.push((port, health_ready_status(port).await)); } if verified.len() != 1 { @@ -967,12 +1015,11 @@ mod tests { } #[test] - fn health_verification_requires_root_kind_and_token_sha() { + fn liveness_verification_requires_root_kind_and_token_sha() { let metadata = metadata(1, Some(8888)); - let health = HealthResponse { - status: Some("healthy".to_string()), + let liveness = DesktopLiveness { + status: Some("alive".to_string()), service: Some("Unsloth UI Backend".to_string()), - version: Some("2026.5.2".to_string()), desktop_protocol_version: Some(1), desktop_manageability_version: Some(1), supports_desktop_auth: Some(true), @@ -983,12 +1030,12 @@ mod tests { token_sha256: Some(token_sha256(TOKEN)), }), }; - assert!(health_verifies_metadata(&health, &metadata)); + assert!(liveness_verifies_metadata(&liveness, &metadata)); - let mut wrong_root = health; + let mut wrong_root = liveness; wrong_root.studio_root_id = Some("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string()); - assert!(!health_verifies_metadata(&wrong_root, &metadata)); + assert!(!liveness_verifies_metadata(&wrong_root, &metadata)); } #[tokio::test] diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 9138f30a58..5a48d26632 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -193,6 +193,9 @@ pub async fn desktop_preflight_result_with_state( if let Some(snapshot) = crate::process::owned_backend_snapshot(state)? { let Some(owner) = snapshot.owner.clone() else { + // TAURI_PORT is emitted only after uvicorn lifespan completes; keep + // this ownerless path on full health so auth/bootstrap are ready. + let probe = match snapshot.port { Some(port) => backend::probe_ownerless_spawned_backend(port).await, None => backend, @@ -494,9 +497,22 @@ mod tests { FakeCli { bin, dir } } + #[cfg(unix)] + fn remove_managed_capability_cache() { + let _ = std::fs::remove_file( + dirs::home_dir() + .unwrap() + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json"), + ); + } + #[cfg(unix)] #[tokio::test] async fn managed_cli_capability_probe_classifies_core_cases() { + remove_managed_capability_cache(); + for (name, script, stale_reason) in [ ( "cap-missing", diff --git a/studio/src-tauri/src/preflight/backend.rs b/studio/src-tauri/src/preflight/backend.rs index 0277ef5149..5a58142667 100644 --- a/studio/src-tauri/src/preflight/backend.rs +++ b/studio/src-tauri/src/preflight/backend.rs @@ -3,7 +3,10 @@ use super::version::{ backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; use serde::{Deserialize, Serialize}; + +use log::info; use std::time::Duration; +use std::time::Instant; #[derive(Debug, Deserialize)] struct DesktopOwnerHealth { @@ -24,6 +27,7 @@ pub(super) struct BackendHealth { } pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Option { + let started = Instant::now(); let url = format!("http://127.0.0.1:{port}/api/health"); let response = client.get(url).send().await.ok()?; if !response.status().is_success() { @@ -40,6 +44,14 @@ pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Optio .and_then(|v| v.as_str()) .map(|s| s == "Unsloth UI Backend") .unwrap_or(false); + info!( + "Desktop preflight: health probe on port {} healthy={} service={} in {}ms", + port, + healthy, + service, + started.elapsed().as_millis() + ); + if !healthy || !service { return None; } diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 51fde51896..57b8365ec5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -2,14 +2,30 @@ use super::types::ManagedProbe; use super::version::{ backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; -use serde::Deserialize; +use log::{info, warn}; +use serde::{Deserialize, Serialize}; +use std::fs; use std::path::{Path, PathBuf}; use std::process::Stdio; -use std::time::Duration; +use std::time::{Duration, Instant, UNIX_EPOCH}; use tokio::io::AsyncReadExt; use tokio::process::Command; -#[derive(Debug, Deserialize)] +const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 2; + +const FNV64_OFFSET_BASIS: u64 = 0xcbf29ce484222325; +const FNV64_PRIME: u64 = 0x100000001b3; +const HASHED_MARKER_MAX_BYTES: u64 = 64 * 1024; + +const FALLBACK_MARKER_NAMES: &[&str] = &[ + "pyvenv.cfg", + "uv.lock", + "requirements.txt", + "python.exe", + "python", +]; + +#[derive(Debug, Clone, Deserialize, Serialize)] struct DesktopCapability { desktop_protocol_version: Option, desktop_manageability_version: Option, @@ -20,7 +36,225 @@ struct DesktopCapability { version: Option, } +#[derive(Debug, Clone, Deserialize, Serialize)] +struct ManagedCapabilityCache { + schema: u16, + bin_path: String, + bin_size: u64, + bin_mtime_ms: u64, + studio_root_id: Option, + marker_path: Option, + marker_size: Option, + marker_mtime_ms: Option, + desktop_protocol_version: u16, + desktop_manageability_version: u16, + capability: DesktopCapability, +} + +#[derive(Debug, Clone)] +struct MarkerFingerprint { + path: String, + size: u64, + mtime_ms: u64, + content_hash: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ManagedBinFingerprint { + bin_path: String, + bin_size: u64, + bin_mtime_ms: u64, + studio_root_id: Option, + marker_path: Option, + marker_size: Option, + marker_mtime_ms: Option, +} + +fn modified_ms(metadata: &fs::Metadata) -> Option { + metadata + .modified() + .ok()? + .duration_since(UNIX_EPOCH) + .ok() + .and_then(|duration| u64::try_from(duration.as_millis()).ok()) +} +fn hash_bytes(hash: u64, bytes: &[u8]) -> u64 { + bytes.iter().fold(hash, |mut next, byte| { + next ^= u64::from(*byte); + next.wrapping_mul(FNV64_PRIME) + }) +} + +fn marker_content_hash(path: &Path, metadata: &fs::Metadata) -> Option { + if metadata.len() > HASHED_MARKER_MAX_BYTES { + return None; + } + fs::read(path) + .ok() + .map(|bytes| hash_bytes(FNV64_OFFSET_BASIS, &bytes)) +} + +fn marker_candidates_for_bin(bin: &Path) -> Vec { + let Some(scripts_dir) = bin.parent() else { + return Vec::new(); + }; + let Some(venv_dir) = scripts_dir.parent() else { + return Vec::new(); + }; + let mut out = Vec::new(); + + #[cfg(unix)] + { + if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) { + for entry in lib_dir.flatten() { + out.push( + entry + .path() + .join("site-packages") + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + ); + } + } + } + for marker_name in FALLBACK_MARKER_NAMES { + out.push(venv_dir.join(marker_name)); + out.push(scripts_dir.join(marker_name)); + } + + out.push( + venv_dir + .join("Lib") + .join("site-packages") + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + ); + out +} + +fn managed_bin_fingerprint(bin: &Path) -> Option { + let bin_metadata = fs::metadata(bin).ok()?; + let bin_path = bin + .canonicalize() + .unwrap_or_else(|_| bin.to_path_buf()) + .to_string_lossy() + .into_owned(); + + let studio_root_id = crate::desktop_backend_owner::read_expected_studio_root_id(); + let mut marker_entries: Vec = marker_candidates_for_bin(bin) + .into_iter() + .filter_map(|path| { + let metadata = fs::metadata(&path).ok()?; + Some(MarkerFingerprint { + path: path + .canonicalize() + .unwrap_or(path.clone()) + .to_string_lossy() + .into_owned(), + size: metadata.len(), + mtime_ms: modified_ms(&metadata)?, + content_hash: marker_content_hash(&path, &metadata), + }) + }) + .collect(); + marker_entries.sort_by(|left, right| left.path.cmp(&right.path)); + let marker_hash = marker_entries + .iter() + .fold(FNV64_OFFSET_BASIS, |hash, marker| { + let next = hash_bytes(hash, marker.path.as_bytes()); + let next = hash_bytes(next, &marker.size.to_le_bytes()); + let next = hash_bytes(next, &marker.mtime_ms.to_le_bytes()); + if let Some(content_hash) = marker.content_hash { + hash_bytes(next, &content_hash.to_le_bytes()) + } else { + next + } + }); + let marker_path = (!marker_entries.is_empty()).then(|| "markers".to_string()); + let marker_size = (!marker_entries.is_empty()).then(|| marker_entries.len() as u64); + let marker_mtime_ms = (!marker_entries.is_empty()).then_some(marker_hash); + + Some(ManagedBinFingerprint { + bin_path, + bin_size: bin_metadata.len(), + bin_mtime_ms: modified_ms(&bin_metadata)?, + studio_root_id, + marker_path, + marker_size, + marker_mtime_ms, + }) +} + +fn capability_cache_path() -> Option { + dirs::home_dir().map(|home| { + home.join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json") + }) +} + +fn cache_matches(cache: &ManagedCapabilityCache, fingerprint: &ManagedBinFingerprint) -> bool { + cache.schema == MANAGED_CAPABILITY_CACHE_SCHEMA + && cache.desktop_protocol_version == DESKTOP_PROTOCOL_VERSION + && cache.desktop_manageability_version == DESKTOP_MANAGEABILITY_VERSION + && cache.bin_path == fingerprint.bin_path + && cache.bin_size == fingerprint.bin_size + && cache.bin_mtime_ms == fingerprint.bin_mtime_ms + && cache.studio_root_id == fingerprint.studio_root_id + && cache.marker_path == fingerprint.marker_path + && cache.marker_size == fingerprint.marker_size + && cache.marker_mtime_ms == fingerprint.marker_mtime_ms + && desktop_capability_ready(&cache.capability) +} + +fn read_cached_capability(fingerprint: &ManagedBinFingerprint) -> Option { + let path = capability_cache_path()?; + let bytes = fs::read(path).ok()?; + let cache = serde_json::from_slice::(&bytes).ok()?; + if cache_matches(&cache, fingerprint) { + Some(cache.capability) + } else { + None + } +} + +fn write_cached_capability(fingerprint: &ManagedBinFingerprint, capability: &DesktopCapability) { + let Some(path) = capability_cache_path() else { + return; + }; + let cache = ManagedCapabilityCache { + schema: MANAGED_CAPABILITY_CACHE_SCHEMA, + bin_path: fingerprint.bin_path.clone(), + bin_size: fingerprint.bin_size, + bin_mtime_ms: fingerprint.bin_mtime_ms, + studio_root_id: fingerprint.studio_root_id.clone(), + marker_path: fingerprint.marker_path.clone(), + marker_size: fingerprint.marker_size, + marker_mtime_ms: fingerprint.marker_mtime_ms, + desktop_protocol_version: DESKTOP_PROTOCOL_VERSION, + desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION, + capability: capability.clone(), + }; + if let Some(parent) = path.parent() { + if fs::create_dir_all(parent).is_err() { + return; + } + } + let Ok(bytes) = serde_json::to_vec_pretty(&cache) else { + return; + }; + if let Err(error) = fs::write(&path, bytes) { + warn!( + "Managed preflight: could not write capability cache: {}", + error + ); + } +} + async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { + let started = Instant::now(); let mut cmd = Command::new(bin); cmd.args(args).stdout(Stdio::null()).stderr(Stdio::null()); @@ -43,20 +277,33 @@ 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, + started.elapsed().as_millis() + ); return false; }; - match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { + let ok = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) => status.success(), _ => { let _ = child.kill().await; let _ = child.wait().await; false } - } + }; + info!( + "Managed preflight probe {:?} finished ok={} in {}ms", + args, + ok, + started.elapsed().as_millis() + ); + ok } async fn probe_cli_capability(bin: &Path) -> Option { + let started = Instant::now(); let mut cmd = Command::new(bin); cmd.args(["studio", "desktop-capabilities", "--json"]) .stdout(Stdio::piped()) @@ -81,6 +328,10 @@ async fn probe_cli_capability(bin: &Path) -> Option { } let Ok(mut child) = cmd.spawn() else { + info!( + "Managed desktop-capabilities probe failed to spawn in {}ms", + started.elapsed().as_millis() + ); return None; }; let Some(mut stdout) = child.stdout.take() else { @@ -92,9 +343,19 @@ async fn probe_cli_capability(bin: &Path) -> Option { Err(_) => { let _ = child.kill().await; let _ = child.wait().await; + info!( + "Managed desktop-capabilities probe timed out in {}ms", + started.elapsed().as_millis() + ); + return None; + } + _ => { + info!( + "Managed desktop-capabilities probe exited unsuccessfully in {}ms", + started.elapsed().as_millis() + ); return None; } - _ => return None, } let mut output = Vec::new(); @@ -102,7 +363,13 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; } - serde_json::from_slice::(&output).ok() + let capability = serde_json::from_slice::(&output).ok(); + info!( + "Managed desktop-capabilities probe finished ok={} in {}ms", + capability.is_some(), + started.elapsed().as_millis() + ); + capability } fn desktop_capability_stale_reason(capability: &DesktopCapability) -> Option { @@ -132,18 +399,48 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { } pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { + let started = Instant::now(); if !run_cli_probe(&bin, &["-h"]).await { + info!( + "Managed preflight: cli unusable for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Stale { bin, reason: "cli_unusable".to_string(), }; } - let capability = probe_cli_capability(&bin).await; - if let Some(capability) = capability { - if desktop_capability_ready(&capability) { + if let Some(fingerprint) = managed_bin_fingerprint(&bin) { + if read_cached_capability(&fingerprint).is_some() { + info!( + "Managed preflight: using cached desktop capability for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Ready { bin }; } + } + + let capability = probe_cli_capability(&bin).await; + if let Some(capability) = capability { + if let Some(fingerprint) = managed_bin_fingerprint(&bin) { + write_cached_capability(&fingerprint, &capability); + } + if desktop_capability_ready(&capability) { + info!( + "Managed preflight: cli ready for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); + return ManagedProbe::Ready { bin }; + } + info!( + "Managed preflight: cli stale for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Stale { bin, reason: desktop_capability_stale_reason(&capability) @@ -151,6 +448,11 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { }; } + info!( + "Managed preflight: desktop capability probe failed for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); ManagedProbe::Stale { bin, reason: "desktop_capability_probe_failed".to_string(), @@ -158,10 +460,17 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { } pub(super) async fn probe_managed_install() -> ManagedProbe { - match crate::process::find_unsloth_binary() { + let started = Instant::now(); + let result = match crate::process::find_unsloth_binary() { Some(bin) => probe_managed_bin(bin).await, None => ManagedProbe::Missing, - } + }; + info!( + "Managed preflight: install probe result {:?} in {}ms", + result, + started.elapsed().as_millis() + ); + result } pub async fn managed_install_ready() -> bool { diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index 01ef77f098..56d9dd2e21 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -735,6 +735,7 @@ pub fn start_backend( } async fn generic_backend_health_ok(port: u16) -> bool { + let started = std::time::Instant::now(); let client = match reqwest::Client::builder() .timeout(Duration::from_secs(2)) .build() @@ -745,49 +746,75 @@ async fn generic_backend_health_ok(port: u16) -> bool { return false; } }; - let response = match client - .get(format!("http://127.0.0.1:{port}/api/health")) - .send() - .await - { - Ok(response) => response, - Err(error) => { + let mut last_status = None; + let mut json = None; + for path in ["/api/liveness", "/api/health"] { + let response = match client + .get(format!("http://127.0.0.1:{port}{path}")) + .send() + .await + { + Ok(response) => response, + Err(error) => { + warn!( + "Backend port candidate {} failed health request: {}", + port, error + ); + return false; + } + }; + if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" { + last_status = Some(response.status()); + continue; + } + if !response.status().is_success() { warn!( - "Backend port candidate {} failed health request: {}", - port, error + "Backend port candidate {} returned HTTP {} from health", + port, + response.status() ); return false; } - }; - if !response.status().is_success() { + json = match response.json::().await { + Ok(json) => Some(json), + Err(error) => { + warn!( + "Backend port candidate {} returned invalid health JSON: {}", + port, error + ); + return false; + } + }; + break; + } + let Some(json) = json else { warn!( "Backend port candidate {} returned HTTP {} from health", port, - response.status() + last_status + .map(|status| status.to_string()) + .unwrap_or_else(|| "unknown".to_string()) ); return false; - } - let json = match response.json::().await { - Ok(json) => json, - Err(error) => { - warn!( - "Backend port candidate {} returned invalid health JSON: {}", - port, error - ); - return false; - } }; - let healthy = json + let live = json .get("status") .and_then(|v| v.as_str()) - .map(|s| s == "healthy") + .map(|s| s == "alive" || s == "healthy") .unwrap_or(false); let service = json .get("service") .and_then(|v| v.as_str()) .map(|s| s == "Unsloth UI Backend") .unwrap_or(false); - healthy && service + info!( + "Backend port candidate {} liveness live={} service={} in {}ms", + port, + live, + service, + started.elapsed().as_millis() + ); + live && service } async fn validate_candidate_port( @@ -798,6 +825,7 @@ async fn validate_candidate_port( generation: u64, port: u16, ) { + let started = std::time::Instant::now(); let owner = { let proc = match state.lock() { Ok(proc) => proc, @@ -852,6 +880,14 @@ async fn validate_candidate_port( } }; + info!( + "Validated backend port candidate {} valid={} emit={} in {}ms", + port, + valid, + should_emit, + started.elapsed().as_millis() + ); + if should_emit { diagnostics::record_backend_port(&diagnostics_state, &session_id, port); info!("Validated backend port: {}", port);