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/internal/model/downloader.go |   68 +++++++++++++++++++++++++++++----
 1 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/privatevoice.src/internal/model/downloader.go b/privatevoice.src/internal/model/downloader.go
index f77c66d..120fe61 100755
--- a/privatevoice.src/internal/model/downloader.go
+++ b/privatevoice.src/internal/model/downloader.go
@@ -1,6 +1,7 @@
 package model
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -27,6 +28,10 @@
 }
 
 func DownloadProfile(profile ModelProfile, urls []string, modelsDir string, progress ProgressCallback) error {
+	return DownloadProfileWithContext(context.Background(), profile, urls, modelsDir, progress)
+}
+
+func DownloadProfileWithContext(ctx context.Context, profile ModelProfile, urls []string, modelsDir string, progress ProgressCallback) error {
 	if err := os.MkdirAll(modelsDir, 0755); err != nil {
 		return fmt.Errorf("failed to create models dir: %w", err)
 	}
@@ -34,10 +39,16 @@
 
 	var lastErr error
 	for i, url := range urls {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		if url == "" {
 			continue
 		}
-		if err := downloadAndInstallFromURL(profile, url, modelsDir, i+1, progress); err != nil {
+		if err := downloadAndInstallFromURL(ctx, profile, url, modelsDir, i+1, progress); err != nil {
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return err
+			}
 			lastErr = err
 			logger.Info("Model install from URL %d failed: %v", i+1, err)
 			continue
@@ -63,7 +74,7 @@
 	}
 }
 
-func downloadAndInstallFromURL(profile ModelProfile, url, modelsDir string, attempt int, progress ProgressCallback) error {
+func downloadAndInstallFromURL(ctx context.Context, profile ModelProfile, url, modelsDir string, attempt int, progress ProgressCallback) error {
 	runID := fmt.Sprintf("%d-%d", time.Now().UnixNano(), attempt)
 	downloadDir := filepath.Join(modelsDir, ".downloads", profile.ID, runID)
 	extractDir := filepath.Join(downloadDir, "extract")
@@ -77,27 +88,42 @@
 	defer os.RemoveAll(stagingDir)
 
 	logger.Info("Downloading model from URL %d: %s", attempt, url)
-	if err := downloadFile(url, archivePath, progress); err != nil {
+	if err := downloadFileWithContext(ctx, url, archivePath, progress); err != nil {
+		return err
+	}
+	if err := ctx.Err(); err != nil {
 		return err
 	}
 
 	logger.Info("Extracting model archive...")
-	if err := extractArchive(archivePath, extractDir); err != nil {
+	if err := extractArchiveWithContext(ctx, archivePath, extractDir); err != nil {
 		return fmt.Errorf("extraction failed: %w", err)
+	}
+	if err := ctx.Err(); err != nil {
+		return err
 	}
 
 	sourceDir, err := findInstallSource(profile, extractDir)
 	if err != nil {
 		return err
 	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 
 	if err := moveOrCopyDir(sourceDir, stagingDir); err != nil {
 		return fmt.Errorf("failed to stage model: %w", err)
+	}
+	if err := ctx.Err(); err != nil {
+		return err
 	}
 
 	validation := ValidateModelDir(profile, stagingDir)
 	if !validation.Valid {
 		return fmt.Errorf("downloaded model is incomplete: missing=%v problems=%v", validation.Missing, validation.Problems)
+	}
+	if err := ctx.Err(); err != nil {
+		return err
 	}
 
 	finalDir := filepath.Join(modelsDir, profile.InstallDirName)
@@ -105,15 +131,25 @@
 }
 
 func downloadFile(url, destPath string, progress ProgressCallback) error {
+	return downloadFileWithContext(context.Background(), url, destPath, progress)
+}
+
+func downloadFileWithContext(ctx context.Context, url, destPath string, progress ProgressCallback) error {
 	const maxAttempts = 6
 	var lastErr error
 	var lastPercent = -1.0
 
 	for attempt := 1; attempt <= maxAttempts; attempt++ {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		downloaded := existingFileSize(destPath)
-		err := downloadFileAttempt(url, destPath, downloaded, &lastPercent, progress)
+		err := downloadFileAttempt(ctx, url, destPath, downloaded, &lastPercent, progress)
 		if err == nil {
 			return nil
+		}
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return err
 		}
 		if isNonRetryableDownloadError(err) {
 			return err
@@ -121,15 +157,19 @@
 		lastErr = err
 		logger.Info("Download attempt %d/%d failed: %v", attempt, maxAttempts, err)
 		if attempt < maxAttempts {
-			time.Sleep(time.Duration(attempt) * time.Second)
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(time.Duration(attempt) * time.Second):
+			}
 		}
 	}
 
 	return fmt.Errorf("download failed after %d attempts: %w", maxAttempts, lastErr)
 }
 
-func downloadFileAttempt(url, destPath string, resumeFrom int64, lastPercent *float64, progress ProgressCallback) error {
-	req, err := http.NewRequest(http.MethodGet, url, nil)
+func downloadFileAttempt(ctx context.Context, url, destPath string, resumeFrom int64, lastPercent *float64, progress ProgressCallback) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
 		return nonRetryableDownloadError{err: err}
 	}
@@ -184,6 +224,9 @@
 	}
 	buf := make([]byte, 32*1024)
 	for {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		n, readErr := resp.Body.Read(buf)
 		if n > 0 {
 			if _, writeErr := out.Write(buf[:n]); writeErr != nil {
@@ -252,9 +295,16 @@
 }
 
 func extractArchive(archivePath, destDir string) error {
-	cmd := exec.Command("tar", "-xf", archivePath, "-C", destDir)
+	return extractArchiveWithContext(context.Background(), archivePath, destDir)
+}
+
+func extractArchiveWithContext(ctx context.Context, archivePath, destDir string) error {
+	cmd := exec.CommandContext(ctx, "tar", "-xf", archivePath, "-C", destDir)
 	output, err := cmd.CombinedOutput()
 	if err != nil {
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return ctxErr
+		}
 		return fmt.Errorf("tar extraction failed: %v, output: %s", err, string(output))
 	}
 	return nil

--
Gitblit v1.9.3