Add Tauri native notifications (unslothai#5273)

Author: wasimysaid

studio/frontend/package.json

@@ -44,6 +44,7 @@
     "@tanstack/react-table": "^8.21.3",
     "@tauri-apps/api": "^2.10.1",
     "@tauri-apps/plugin-clipboard-manager": "^2.3.2",
+    "@tauri-apps/plugin-notification": "^2.3.3",
     "@tauri-apps/plugin-opener": "^2.5.3",
     "@tauri-apps/plugin-process": "^2.3.1",
     "@tauri-apps/plugin-updater": "^2.10.1",

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -4,6 +4,12 @@
 import { createElement, useCallback, useRef, useState } from "react";
 import { toast } from "sonner";
 import { consumeNativePathToken } from "@/features/native-intents/api";
+import {
+  notifyNative,
+  primeNativeNotificationPermission,
+  safeNotificationLabel,
+  sanitizeNotificationBody,
+} from "@/lib/native-notifications";
 import { ModelLoadDescription } from "../components/model-load-status";
 import {
   getDownloadProgress,
@@ -168,6 +174,7 @@ export function useChatModelRuntime() {
   const loadAbortRef = useRef<AbortController | null>(null);
   const loadingModelRef = useRef<typeof loadingModel>(null);
   const loadToastIdRef = useRef<string | number | null>(null);
+  const loadAttemptRef = useRef(0);
   const loadToastDismissedRef = useRef(false);
   const cancelUnloadPendingRef = useRef(false);
 
@@ -369,6 +376,10 @@ export function useChatModelRuntime() {
       const isLora =
         explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false;
       const displayName = model?.name || lora?.name || modelId;
+      const loadAttemptId = ++loadAttemptRef.current;
+      primeNativeNotificationPermission().catch(() => undefined);
+      const notificationModelKey = `${modelId}:${ggufVariant ?? ""}:${loadAttemptId}`;
+      const safeModelName = safeNotificationLabel(displayName, "The model");
       const currentCheckpoint =
         useChatRuntimeStore.getState().params.checkpoint;
       const previousCheckpoint = currentCheckpoint;
@@ -813,6 +824,12 @@ export function useChatModelRuntime() {
                   },
                 });
               }
+              notifyNative({
+                key: `model-downloaded:${notificationModelKey}`,
+                title: "Model downloaded",
+                body: `${safeModelName} finished downloading and is loading into memory.`,
+                requestPermission: false,
+              }).catch(() => undefined);
               // Keep polling: the mmap branch below takes over from here.
             }
           } catch {
@@ -898,6 +915,12 @@ export function useChatModelRuntime() {
               duration: 2000,
             });
           }
+          notifyNative({
+            key: `model-loaded:${notificationModelKey}`,
+            title: "Model ready",
+            body: `${safeModelName} is loaded and ready to chat.`,
+            requestPermission: false,
+          }).catch(() => undefined);
         } catch (err) {
           if (!abortCtrl.signal.aborted) {
             const message =
@@ -912,6 +935,12 @@ export function useChatModelRuntime() {
                 duration: 5000,
               });
             }
+            notifyNative({
+              key: `model-load-failed:${notificationModelKey}`,
+              title: "Model failed to load",
+              body: sanitizeNotificationBody(message, "The model failed to load."),
+              requestPermission: false,
+            }).catch(() => undefined);
           }
           throw err;
         } finally {

studio/frontend/src/features/training/hooks/use-training-actions.ts

@@ -1,19 +1,20 @@
 // 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 { primeNativeNotificationPermission } from "@/lib/native-notifications";
 import { useCallback } from "react";
+import { toast } from "sonner";
 import { checkDatasetFormat } from "../api/datasets-api";
 import { getTrainingRun } from "../api/history-api";
 import { buildTrainingStartPayload } from "../api/mappers";
-import { startTraining, stopTraining, resetTraining } from "../api/train-api";
+import { resetTraining, startTraining, stopTraining } from "../api/train-api";
 import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
 import { validateTrainingConfig } from "../lib/validation";
 import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store";
 import { useTrainingConfigStore } from "../stores/training-config-store";
 import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
 import type { TrainingStartRequest } from "../types/api";
 import type { TrainingConfigState } from "../types/config";
-import { toast } from "sonner";
 
 /** Chatml → format-specific role remap (only for formats that differ from chatml). */
 const ROLE_REMAP: Record<string, Record<string, string>> = {
@@ -50,6 +51,8 @@ export function useTrainingActions() {
       return false;
     }
 
+    primeNativeNotificationPermission().catch(() => undefined);
+
     runtimeStore.setStartResources(
       config.selectedModel ?? null,
       getHfDatasetName(config),
@@ -174,6 +177,8 @@ export function useTrainingActions() {
         throw new Error("Only stopped runs with a saved checkpoint can be resumed.");
       }
 
+      primeNativeNotificationPermission().catch(() => undefined);
+
       const config = useTrainingConfigStore.getState();
       const savedConfig = detail.config as Partial<TrainingStartRequest>;
       const payload = {

studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts

@@ -2,25 +2,98 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { hasAuthToken } from "@/features/auth";
+import {
+  notifyNative,
+  safeNotificationLabel,
+  sanitizeNotificationBody,
+} from "@/lib/native-notifications";
 import { useEffect } from "react";
 import {
   getTrainingMetrics,
   getTrainingStatus,
   isAbortError,
   streamTrainingProgress,
 } from "../api/train-api";
+import { useTrainingConfigStore } from "../stores/training-config-store";
 import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
 import type { TrainingRuntimeStore } from "../types/runtime";
 
 const STATUS_POLL_INTERVAL_MS = 3000;
 const METRICS_POLL_INTERVAL_MS = 5000;
 const IDLE_POLL_INTERVAL_MS = 30000;
 const STREAM_RECONNECT_DELAY_MS = 1500;
+const GENERIC_TRAINING_STREAM_ERROR = "Training stream error";
+const DEFAULT_TRAINING_ERROR_BODY = "Your training run stopped with an error.";
 
 function shouldUseLiveSync(state: TrainingRuntimeStore): boolean {
   return state.isTrainingRunning || state.phase === "training";
 }
 
+function getTrainingErrorBody(state: TrainingRuntimeStore): string {
+  return sanitizeNotificationBody(
+    state.error ?? state.message,
+    DEFAULT_TRAINING_ERROR_BODY,
+  );
+}
+
+function shouldNotifyTrainingError(
+  before: TrainingRuntimeStore,
+  after: TrainingRuntimeStore,
+): boolean {
+  if (after.phase !== "error") {
+    return false;
+  }
+  if (before.phase !== "error" || before.jobId !== after.jobId) {
+    return true;
+  }
+  if (before.error !== GENERIC_TRAINING_STREAM_ERROR) {
+    return false;
+  }
+
+  return getTrainingErrorBody(before) !== getTrainingErrorBody(after);
+}
+
+function maybeNotifyTrainingTerminalTransition(
+  before: TrainingRuntimeStore,
+  after: TrainingRuntimeStore,
+): void {
+  if (!after.jobId) {
+    return;
+  }
+  if (
+    before.stopRequested ||
+    after.stopRequested ||
+    after.phase === "stopped"
+  ) {
+    return;
+  }
+
+  const sameJob = before.jobId === after.jobId;
+  const samePhase = before.phase === after.phase;
+
+  if (after.phase === "completed" && !(sameJob && samePhase)) {
+    const model = safeNotificationLabel(
+      after.startModelName ?? useTrainingConfigStore.getState().selectedModel,
+      "Your training run",
+    );
+    notifyNative({
+      key: `training-completed:${after.jobId}`,
+      title: "Training finished",
+      body: `${model} is complete.`,
+      requestPermission: false,
+    }).catch(() => undefined);
+  }
+
+  if (shouldNotifyTrainingError(before, after)) {
+    notifyNative({
+      key: `training-error:${after.jobId}`,
+      title: "Training failed",
+      body: getTrainingErrorBody(after),
+      requestPermission: false,
+    }).catch(() => undefined);
+  }
+}
+
 export function useTrainingRuntimeLifecycle(): void {
   useEffect(() => {
     let disposed = false;
@@ -62,7 +135,9 @@ export function useTrainingRuntimeLifecycle(): void {
       }
     };
 
-    const pollStatus = async () => {
+    const pollStatus = async (options?: {
+      suppressNativeNotifications?: boolean;
+    }) => {
       if (!hasAuthToken()) return;
       const gen = runtimeStore.getState().resetGeneration;
       try {
@@ -71,9 +146,13 @@ export function useTrainingRuntimeLifecycle(): void {
           return;
         }
 
+        const previousState = runtimeStore.getState();
         runtimeStore.getState().applyStatus(status);
 
         const nextState = runtimeStore.getState();
+        if (!options?.suppressNativeNotifications) {
+          maybeNotifyTrainingTerminalTransition(previousState, nextState);
+        }
         if (shouldUseLiveSync(nextState)) {
           void ensureStream();
         } else {
@@ -124,7 +203,7 @@ export function useTrainingRuntimeLifecycle(): void {
             }
 
             if (event.event === "error") {
-              liveStore.setRuntimeError("Training stream error");
+              liveStore.setRuntimeError(GENERIC_TRAINING_STREAM_ERROR);
               stopStream();
             }
           },
@@ -154,7 +233,10 @@ export function useTrainingRuntimeLifecycle(): void {
     const hydrate = async () => {
       runtimeStore.getState().setHydrating(true);
       try {
-        await Promise.all([pollStatus(), pollMetrics()]);
+        await Promise.all([
+          pollStatus({ suppressNativeNotifications: true }),
+          pollMetrics(),
+        ]);
       } finally {
         if (!disposed) {
           runtimeStore.getState().setHydrating(false);
@@ -172,7 +254,12 @@ export function useTrainingRuntimeLifecycle(): void {
 
     const statusTimer = setInterval(() => {
       const s = runtimeStore.getState();
-      if (isIdle() && s.hasHydrated) return;
+      if (!s.hasHydrated || s.isHydrating) {
+        return;
+      }
+      if (isIdle()) {
+        return;
+      }
       void pollStatus();
     }, STATUS_POLL_INTERVAL_MS);
 
@@ -187,7 +274,13 @@ export function useTrainingRuntimeLifecycle(): void {
     // Low-frequency background poll to recover from failed hydration or detect
     // out-of-band state changes (e.g. training started from another client).
     const idleTimer = setInterval(() => {
-      if (!isIdle()) return;
+      const s = runtimeStore.getState();
+      if (!s.hasHydrated || s.isHydrating) {
+        return;
+      }
+      if (!isIdle()) {
+        return;
+      }
       void pollStatus();
     }, IDLE_POLL_INTERVAL_MS);