From 7999e66c9c78ada3666eb7cea7c22a352bb4cbf0 Mon Sep 17 00:00:00 2001
From: Ariver <shanghai3168@gmail.com>
Date: Wed, 24 Jun 2026 02:30:51 +0800
Subject: [PATCH] Add cancellable model downloads

---
 privatevoice.src/services/engine_service.go |   83 ++++++++++++++++++++++++++++++++++++-----
 1 files changed, 73 insertions(+), 10 deletions(-)

diff --git a/privatevoice.src/services/engine_service.go b/privatevoice.src/services/engine_service.go
index b6c00b6..4dbfeb3 100755
--- a/privatevoice.src/services/engine_service.go
+++ b/privatevoice.src/services/engine_service.go
@@ -1,6 +1,7 @@
 package services
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"voicesnap/internal/config"
@@ -15,13 +16,16 @@
 
 // EngineService provides engine status and model management to the frontend.
 type EngineService struct {
-	cfg          *config.Config
-	app          *application.App
-	initCallback func()
-	mu           sync.RWMutex
-	status       string
-	hardwareInfo string
-	lastError    string
+	cfg            *config.Config
+	app            *application.App
+	initCallback   func()
+	mu             sync.RWMutex
+	status         string
+	hardwareInfo   string
+	lastError      string
+	downloadMu     sync.Mutex
+	downloadID     string
+	downloadCancel context.CancelFunc
 }
 
 func NewEngineService(cfg *config.Config) *EngineService {
@@ -106,7 +110,13 @@
 		return fmt.Errorf("no download URL configured for model %s", profile.ID)
 	}
 
-	err = model.DownloadProfile(profile, profile.DownloadURLs, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
+	ctx, finish, err := s.beginModelDownload(profile.ID)
+	if err != nil {
+		return err
+	}
+	defer finish()
+
+	err = model.DownloadProfileWithContext(ctx, profile, profile.DownloadURLs, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
 		if s.app != nil {
 			s.app.Event.Emit("model:download-progress", map[string]interface{}{
 				"percent":    percent,
@@ -126,6 +136,24 @@
 	config.Save(s.cfg)
 	s.ReloadCurrentModel()
 	return nil
+}
+
+func (s *EngineService) CancelModelDownload(modelID string) bool {
+	s.downloadMu.Lock()
+	activeID := s.downloadID
+	if s.downloadCancel == nil || (modelID != "" && activeID != modelID) {
+		s.downloadMu.Unlock()
+		return false
+	}
+	s.downloadCancel()
+	s.downloadMu.Unlock()
+
+	if s.app != nil {
+		s.app.Event.Emit("model:download-cancelled", map[string]interface{}{
+			"modelID": activeID,
+		})
+	}
+	return true
 }
 
 func (s *EngineService) modelStatusMap(profile model.ModelProfile, current modelselection.CurrentModel) map[string]interface{} {
@@ -179,12 +207,21 @@
 
 // DownloadModel downloads the ASR model with progress events.
 func (s *EngineService) DownloadModel(primaryURL, fallbackURL string) error {
-	err := model.Download(primaryURL, fallbackURL, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
+	profile := model.DefaultModelProfile()
+	ctx, finish, err := s.beginModelDownload(profile.ID)
+	if err != nil {
+		return err
+	}
+	defer finish()
+
+	err = model.DownloadProfileWithContext(ctx, profile, []string{primaryURL, fallbackURL}, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
 		if s.app != nil {
 			s.app.Event.Emit("model:download-progress", map[string]interface{}{
 				"percent":    percent,
 				"downloaded": downloaded,
 				"total":      total,
+				"modelID":    profile.ID,
+				"modelName":  profile.DisplayName,
 			})
 		}
 	})
@@ -204,7 +241,13 @@
 		return fmt.Errorf("no download URL configured for model %s", current.ModelID)
 	}
 
-	err := model.DownloadProfile(current.Profile, current.Profile.DownloadURLs, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
+	ctx, finish, err := s.beginModelDownload(current.ModelID)
+	if err != nil {
+		return err
+	}
+	defer finish()
+
+	err = model.DownloadProfileWithContext(ctx, current.Profile, current.Profile.DownloadURLs, paths.ModelsRoot(), func(percent float64, downloaded, total int64) {
 		if s.app != nil {
 			s.app.Event.Emit("model:download-progress", map[string]interface{}{
 				"percent":    percent,
@@ -252,6 +295,26 @@
 	return modelselection.Resolve(s.cfg, language.NewSystemDetector())
 }
 
+func (s *EngineService) beginModelDownload(modelID string) (context.Context, func(), error) {
+	s.downloadMu.Lock()
+	defer s.downloadMu.Unlock()
+	if s.downloadCancel != nil {
+		return nil, nil, fmt.Errorf("model %s is already downloading", s.downloadID)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	s.downloadID = modelID
+	s.downloadCancel = cancel
+	finish := func() {
+		s.downloadMu.Lock()
+		if s.downloadID == modelID {
+			s.downloadID = ""
+			s.downloadCancel = nil
+		}
+		s.downloadMu.Unlock()
+	}
+	return ctx, finish, nil
+}
+
 func (s *EngineService) allowedModelProfile(modelID string) (model.ModelProfile, error) {
 	profile, err := model.GetModelProfile(modelID)
 	if err != nil {

--
Gitblit v1.9.3