Harden model download recovery

Sources/AppDelegate.swift

@@ -58,17 +58,17 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
         let current = UserDefaults.standard.string(forKey: Transcriber.modelKey) ?? Transcriber.defaultModel
         let modelMenu = NSMenu()
 
-        tinyModelItem = NSMenuItem(title: "Tiny (faster)", action: #selector(selectModel(_:)), keyEquivalent: "")
+        tinyModelItem = NSMenuItem(title: modelMenuTitle(file: "ggml-tiny.en.bin"), action: #selector(selectModel(_:)), keyEquivalent: "")
         tinyModelItem.representedObject = "ggml-tiny.en.bin"
         tinyModelItem.state = "ggml-tiny.en.bin" == current ? .on : .off
         modelMenu.addItem(tinyModelItem)
 
-        baseModelItem = NSMenuItem(title: downloadableModelTitle(file: "ggml-base.en.bin", label: "Base"), action: #selector(selectModel(_:)), keyEquivalent: "")
+        baseModelItem = NSMenuItem(title: modelMenuTitle(file: "ggml-base.en.bin"), action: #selector(selectModel(_:)), keyEquivalent: "")
         baseModelItem.representedObject = "ggml-base.en.bin"
         baseModelItem.state = "ggml-base.en.bin" == current ? .on : .off
         modelMenu.addItem(baseModelItem)
 
-        smallModelItem = NSMenuItem(title: downloadableModelTitle(file: "ggml-small.en.bin", label: "Small (~466MB)"), action: #selector(selectModel(_:)), keyEquivalent: "")
+        smallModelItem = NSMenuItem(title: modelMenuTitle(file: "ggml-small.en.bin"), action: #selector(selectModel(_:)), keyEquivalent: "")
         smallModelItem.representedObject = "ggml-small.en.bin"
         smallModelItem.state = "ggml-small.en.bin" == current ? .on : .off
         modelMenu.addItem(smallModelItem)
@@ -94,6 +94,7 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
     func menuWillOpen(_ menu: NSMenu) {
         guard menu === statusItem.menu else { return }
         refreshPermissionItems()
+        refreshModelMenu()
     }
 
     private func refreshPermissionItems() {
@@ -150,10 +151,37 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
 
     // MARK: - Model Selection
 
+    private func baseModelLabel(for file: String) -> String {
+        switch file {
+        case "ggml-tiny.en.bin":
+            return "Tiny (faster)"
+        case "ggml-base.en.bin":
+            return "Base"
+        case "ggml-small.en.bin":
+            return "Small (~466MB)"
+        default:
+            return file
+        }
+    }
+
+    private func userModelDirectoryURL() -> URL {
+        FileManager.default.homeDirectoryForCurrentUser
+            .appendingPathComponent(".yell", isDirectory: true)
+            .appendingPathComponent("models", isDirectory: true)
+    }
+
+    private func userModelURL(for file: String) -> URL {
+        userModelDirectoryURL().appendingPathComponent(file)
+    }
+
+    private func downloadedModelPath(for file: String) -> String? {
+        guard Bundle.main.path(forResource: file, ofType: nil) == nil else { return nil }
+        let path = userModelURL(for: file).path
+        return FileManager.default.fileExists(atPath: path) ? path : nil
+    }
+
     private func modelExists(_ file: String) -> Bool {
-        if Bundle.main.path(forResource: file, ofType: nil) != nil { return true }
-        let home = FileManager.default.homeDirectoryForCurrentUser.path
-        return FileManager.default.fileExists(atPath: "\(home)/.yell/models/\(file)")
+        Bundle.main.path(forResource: file, ofType: nil) != nil || downloadedModelPath(for: file) != nil
     }
 
     private var currentModelFile: String {
@@ -168,10 +196,14 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
         modelExists(file) ? label : "\(label) — Download"
     }
 
+    private func modelMenuTitle(file: String) -> String {
+        downloadableModelTitle(file: file, label: baseModelLabel(for: file))
+    }
+
     private func refreshModelMenu(selectedFile: String? = nil, useCurrentSelection: Bool = true) {
-        tinyModelItem.title = "Tiny (faster)"
-        baseModelItem.title = downloadableModelTitle(file: "ggml-base.en.bin", label: "Base")
-        smallModelItem.title = downloadableModelTitle(file: "ggml-small.en.bin", label: "Small (~466MB)")
+        tinyModelItem.title = modelMenuTitle(file: "ggml-tiny.en.bin")
+        baseModelItem.title = modelMenuTitle(file: "ggml-base.en.bin")
+        smallModelItem.title = modelMenuTitle(file: "ggml-small.en.bin")
 
         let file = useCurrentSelection ? (selectedFile ?? currentModelFile) : selectedFile
         modelItems.forEach { item in
@@ -191,6 +223,58 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
         refreshModelMenu(selectedFile: selectedFile, useCurrentSelection: useCurrentSelection)
     }
 
+    private func installDownloadedModel(from temporaryURL: URL, to destinationURL: URL) throws {
+        let fileManager = FileManager.default
+        try fileManager.createDirectory(at: userModelDirectoryURL(), withIntermediateDirectories: true)
+        if fileManager.fileExists(atPath: destinationURL.path) {
+            try fileManager.removeItem(at: destinationURL)
+        }
+        try fileManager.moveItem(at: temporaryURL, to: destinationURL)
+    }
+
+    private func stagedDownloadedModelURL(for file: String) -> URL {
+        userModelDirectoryURL().appendingPathComponent("\(file).download")
+    }
+
+    private func persistDownloadedModelTemporarily(from temporaryURL: URL, for file: String) throws -> URL {
+        let fileManager = FileManager.default
+        let stagedURL = stagedDownloadedModelURL(for: file)
+        try fileManager.createDirectory(at: userModelDirectoryURL(), withIntermediateDirectories: true)
+        if fileManager.fileExists(atPath: stagedURL.path) {
+            try fileManager.removeItem(at: stagedURL)
+        }
+        try fileManager.moveItem(at: temporaryURL, to: stagedURL)
+        return stagedURL
+    }
+
+    private func removeDownloadedModelTemporarily(at stagedURL: URL) {
+        try? FileManager.default.removeItem(at: stagedURL)
+    }
+
+    private func discardInvalidDownloadedModel(file: String) {
+        guard let path = downloadedModelPath(for: file) else { return }
+        try? FileManager.default.removeItem(atPath: path)
+        refreshModelMenu()
+    }
+
+    private func showInvalidDownloadedModelAlert(file: String) {
+        let alert = NSAlert()
+        alert.messageText = "Downloaded Model Is Invalid"
+        alert.informativeText = "\(file) could not be loaded and was removed from ~/.yell/models/. Download it again to use it."
+        alert.alertStyle = .warning
+        alert.addButton(withTitle: "OK")
+        alert.runModal()
+    }
+
+    private func showModelLoadFailedAlert(file: String) {
+        let alert = NSAlert()
+        alert.messageText = "Couldn't Load \(file)"
+        alert.informativeText = "Yell restored the previous model. Try downloading \(file) again if the problem persists."
+        alert.alertStyle = .warning
+        alert.addButton(withTitle: "OK")
+        alert.runModal()
+    }
+
     @objc private func selectModel(_ sender: NSMenuItem) {
         guard !isSwitchingModel else { return }
         guard let file = sender.representedObject as? String else { return }
@@ -208,11 +292,18 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
             return
         }
 
+        if let path = downloadedModelPath(for: file),
+           !Transcriber.canLoadModel(atPath: path) {
+            discardInvalidDownloadedModel(file: file)
+            showInvalidDownloadedModelAlert(file: file)
+            return
+        }
+
         switchModel(to: file, selecting: sender, previousFile: previousFile)
     }
 
     private func switchModel(to file: String, selecting item: NSMenuItem, previousFile: String, beginOperation: Bool = true) {
-        let baseLabel = item.title.components(separatedBy: " —").first ?? item.title
+        let baseLabel = baseModelLabel(for: file)
         if beginOperation {
             beginModelOperation(statusTitle: "\(baseLabel) — Loading…", on: item)
         }
@@ -232,13 +323,18 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
         UserDefaults.standard.set(previousFile, forKey: Transcriber.modelKey)
         transcriber.reload { [weak self] restored in
             guard let self else { return }
-            self.showModelSwitchFailedAlert(file: attemptedFile, restoredPreviousModel: restored)
-            self.endModelOperation(selectedFile: restored ? previousFile : nil, useCurrentSelection: restored)
+            guard restored else {
+                self.endModelOperation(selectedFile: nil, useCurrentSelection: false)
+                self.showModelMissingAlert()
+                return
+            }
+            self.showModelLoadFailedAlert(file: attemptedFile)
+            self.endModelOperation(selectedFile: previousFile)
         }
     }
 
     private func downloadModel(file: String, menuItem: NSMenuItem, previousFile: String) {
-        let baseLabel = menuItem.title.components(separatedBy: " —").first ?? menuItem.title
+        let baseLabel = baseModelLabel(for: file)
         beginModelOperation(statusTitle: "\(baseLabel) — Downloading…", on: menuItem)
 
         guard let url = URL(string: "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/\(file)") else {
@@ -249,31 +345,48 @@ class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
 
         URLSession.shared.downloadTask(with: url) { [weak self] tmpURL, response, error in
             guard let self else { return }
-            DispatchQueue.main.async {
-                let httpStatus = (response as? HTTPURLResponse)?.statusCode
-                guard error == nil, let tmpURL, httpStatus == 200 else {
+            let httpStatus = (response as? HTTPURLResponse)?.statusCode
+            guard error == nil, let tmpURL, httpStatus == 200 else {
+                DispatchQueue.main.async {
                     self.endModelOperation(selectedFile: previousFile)
                     let details = httpStatus.map { "Server responded with HTTP \($0)." }
                     self.showModelDownloadFailedAlert(file: file, details: details)
-                    return
                 }
+                return
+            }
 
-                let home = FileManager.default.homeDirectoryForCurrentUser.path
-                let destDir = "\(home)/.yell/models"
-                let dest = URL(fileURLWithPath: "\(destDir)/\(file)")
-
-                do {
-                    try FileManager.default.createDirectory(atPath: destDir, withIntermediateDirectories: true)
-                    if FileManager.default.fileExists(atPath: dest.path) {
-                        try FileManager.default.removeItem(at: dest)
-                    }
-                    try FileManager.default.moveItem(at: tmpURL, to: dest)
-                } catch {
+            let stagedURL: URL
+            do {
+                stagedURL = try self.persistDownloadedModelTemporarily(from: tmpURL, for: file)
+            } catch {
+                DispatchQueue.main.async {
                     self.endModelOperation(selectedFile: previousFile)
                     self.showModelDownloadFailedAlert(file: file, details: error.localizedDescription)
-                    return
                 }
+                return
+            }
 
+            guard Transcriber.canLoadModel(atPath: stagedURL.path) else {
+                self.removeDownloadedModelTemporarily(at: stagedURL)
+                DispatchQueue.main.async {
+                    self.endModelOperation(selectedFile: previousFile)
+                    self.showInvalidDownloadedModelAlert(file: file)
+                }
+                return
+            }
+
+            do {
+                try self.installDownloadedModel(from: stagedURL, to: self.userModelURL(for: file))
+            } catch {
+                self.removeDownloadedModelTemporarily(at: stagedURL)
+                DispatchQueue.main.async {
+                    self.endModelOperation(selectedFile: previousFile)
+                    self.showModelDownloadFailedAlert(file: file, details: error.localizedDescription)
+                }
+                return
+            }
+
+            DispatchQueue.main.async {
                 self.switchModel(to: file, selecting: menuItem, previousFile: previousFile, beginOperation: false)
             }
         }.resume()

Sources/Transcriber.swift

@@ -49,6 +49,20 @@ class Transcriber {
         }
     }
 
+    private static func createContext(modelPath: String) -> OpaquePointer? {
+        let cparams = whisper_context_default_params()
+        return whisper_init_from_file_with_params(modelPath, cparams)
+    }
+
+    static func canLoadModel(atPath path: String) -> Bool {
+        guard FileManager.default.fileExists(atPath: path),
+              let ctx = createContext(modelPath: path) else {
+            return false
+        }
+        whisper_free(ctx)
+        return true
+    }
+
     private func reloadLocked() -> Bool {
         if let ctx = context {
             whisper_free(ctx)
@@ -63,8 +77,7 @@ class Transcriber {
             return false
         }
 
-        let cparams = whisper_context_default_params()
-        context = whisper_init_from_file_with_params(modelPath, cparams)
+        context = Transcriber.createContext(modelPath: modelPath)
         if context == nil {
             print("Failed to initialize whisper context")
             return false

build.sh

@@ -1,6 +1,9 @@
 #!/bin/bash
 set -euo pipefail
 
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "$SCRIPT_DIR/model-checksums.sh"
+
 APP_NAME="Yell"
 BUILD_DIR="build"
 WHISPER_DIR="vendor/whisper.cpp"
@@ -116,6 +119,10 @@ cp "$BUILD_DIR/$APP_NAME" "$APP_BUNDLE/Contents/MacOS/$APP_NAME"
 # Bundle tiny.en model
 TINY_MODEL="$HOME/.yell/models/ggml-tiny.en.bin"
 if [ -f "$TINY_MODEL" ]; then
+    if ! model_checksum_matches "$TINY_MODEL" "ggml-tiny.en.bin"; then
+        echo "Invalid ggml-tiny.en.bin checksum at $TINY_MODEL — run ./download-model.sh ggml-tiny.en.bin" >&2
+        exit 1
+    fi
     echo "Bundling ggml-tiny.en.bin..."
     cp "$TINY_MODEL" "$APP_BUNDLE/Contents/Resources/ggml-tiny.en.bin"
 else

download-model.sh

@@ -1,6 +1,9 @@
 #!/bin/bash
 set -euo pipefail
 
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "$SCRIPT_DIR/model-checksums.sh"
+
 MODEL_DIR="$HOME/.yell/models"
 BASE_URL="https://huggingface.co/ggerganov/whisper.cpp/resolve/main"
 DEFAULT_MODELS=(
@@ -14,13 +17,43 @@ mkdir -p "$MODEL_DIR"
 download_if_missing() {
     local name="$1"
     local path="$MODEL_DIR/$name"
+    local expected_checksum=""
+
+    if ! expected_checksum="$(model_checksum_for "$name")"; then
+        if [ -f "$path" ]; then
+            echo "No checksum is defined for $name, keeping the existing file." >&2
+            return
+        fi
+
+        echo "No checksum is defined for $name" >&2
+        return 1
+    fi
+
     if [ -f "$path" ]; then
-        echo "$name already exists, skipping."
-    else
-        echo "Downloading $name..."
-        curl -fL --retry 3 --progress-bar -o "$path" "$BASE_URL/$name"
-        echo "Downloaded to $path"
+        if model_checksum_matches "$path" "$name"; then
+            echo "$name already exists, skipping."
+            return
+        fi
+
+        echo "$name exists but failed checksum validation, re-downloading..." >&2
+        rm -f "$path"
+    fi
+
+    echo "Downloading $name..."
+    local tmp_path="$path.download"
+    rm -f "$tmp_path"
+    local curl_progress="--no-progress-meter"
+    if [ -t 2 ]; then
+        curl_progress="--progress-bar"
+    fi
+    curl -fL --retry 3 "$curl_progress" -o "$tmp_path" "$BASE_URL/$name"
+    if ! model_checksum_matches "$tmp_path" "$name"; then
+        rm -f "$tmp_path"
+        echo "$name failed checksum validation after download (expected $expected_checksum)." >&2
+        return 1
     fi
+    mv "$tmp_path" "$path"
+    echo "Downloaded to $path"
 }
 
 MODELS=("$@")

model-checksums.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+model_checksum_for() {
+    case "$1" in
+        ggml-tiny.en.bin)
+            echo "921e4cf8686fdd993dcd081a5da5b6c365bfde1162e72b08d75ac75289920b1f"
+            ;;
+        ggml-base.en.bin)
+            echo "a03779c86df3323075f5e796cb2ce5029f00ec8869eee3fdfb897afe36c6d002"
+            ;;
+        ggml-small.en.bin)
+            echo "c6138d6d58ecc8322097e0f987c32f1be8bb0a18532a3f88f734d1bbf9c41e5d"
+            ;;
+        *)
+            return 1
+            ;;
+    esac
+}
+
+model_checksum_matches() {
+    local path="$1"
+    local name="$2"
+    local expected actual
+
+    expected="$(model_checksum_for "$name")" || return 1
+    actual="$(shasum -a 256 "$path" | awk '{print $1}')"
+    [ "$actual" = "$expected" ]
+}