From 5a5d3260f40542584acaeb2d0aa9a126e43466af Mon Sep 17 00:00:00 2001
From: Eva Ho <hoyyeva@gmail.com>
Date: Wed, 17 Dec 2025 16:57:43 -0500
Subject: [PATCH] fix behaviour when switching between enabled and disabled

---
 app/ui/ui.go           | 41 ++++++++++++++++++++++++++++++++---------
 app/updater/updater.go | 34 ++++++++++++++++++++++++++++++----
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/app/ui/ui.go b/app/ui/ui.go
index a5d27739e..f75ec1b43 100644
--- a/app/ui/ui.go
+++ b/app/ui/ui.go
@@ -119,6 +119,7 @@ type UpdaterInterface interface {
 	CheckForUpdate(ctx context.Context) (bool, string, error)
 	DownloadUpdate(ctx context.Context, updateVersion string) error
 	InstallAndRestart() error
+	CancelOngoingDownload()
 }
 
 func (s *Server) log() *slog.Logger {
@@ -1466,15 +1467,41 @@ func (s *Server) settings(w http.ResponseWriter, r *http.Request) error {
 
 	// Update tray notification based on auto-update toggle
 	if old.AutoUpdateEnabled && !settings.AutoUpdateEnabled {
-		// Auto-update disabled: clear tray notification
+		// Auto-update disabled: cancel any ongoing download and clear tray notification
+		if s.Updater != nil {
+			s.Updater.CancelOngoingDownload()
+		}
 		if s.ClearUpdateAvailableFunc != nil {
 			s.ClearUpdateAvailableFunc()
 		}
 	} else if !old.AutoUpdateEnabled && settings.AutoUpdateEnabled {
-		// Auto-update enabled: show tray notification if update is pending
-		if (updater.IsUpdatePending() || updater.UpdateDownloaded) && s.UpdateAvailableFunc != nil {
-			s.UpdateAvailableFunc()
-		}
+		// Auto-update enabled: check for updates and download if available
+		go func() {
+			// First, show notification if update is already pending
+			if (updater.IsUpdatePending() || updater.UpdateDownloaded) && s.UpdateAvailableFunc != nil {
+				s.UpdateAvailableFunc()
+			} else if s.Updater != nil {
+				// Otherwise, immediately check for and download new updates
+				slog.Info("auto-update re-enabled, checking for updates")
+				available, updateVersion, err := s.Updater.CheckForUpdate(r.Context())
+				if err != nil {
+					slog.Error("failed to check for update after re-enabling auto-update", "error", err)
+					return
+				}
+				if available {
+					slog.Info("update available, starting download", "version", updateVersion)
+					err := s.Updater.DownloadUpdate(r.Context(), updateVersion)
+					if err != nil {
+						slog.Error("failed to download update", "error", err)
+						return
+					}
+					// Show tray notification after successful download
+					if s.UpdateAvailableFunc != nil {
+						s.UpdateAvailableFunc()
+					}
+				}
+			}
+		}()
 	}
 
 	if old.ContextLength != settings.ContextLength ||
@@ -1631,10 +1658,6 @@ func (s *Server) installUpdate(w http.ResponseWriter, r *http.Request) error {
 		return fmt.Errorf("no update downloaded")
 	}
 
-	// Check if we can actually upgrade (not in dev mode)
-	if updater.BundlePath == "" {
-		return fmt.Errorf("cannot install updates in development mode")
-	}
 
 	// Send response before restarting
 	response := map[string]any{
diff --git a/app/updater/updater.go b/app/updater/updater.go
index 7b03aa841..e40916dde 100644
--- a/app/updater/updater.go
+++ b/app/updater/updater.go
@@ -19,6 +19,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ollama/ollama/app/store"
@@ -134,15 +135,27 @@ func (u *Updater) checkForUpdate(ctx context.Context) (bool, UpdateResponse) {
 }
 
 func (u *Updater) DownloadNewRelease(ctx context.Context, updateResp UpdateResponse) error {
+	// Create a cancellable context for this download
+	downloadCtx, cancel := context.WithCancel(ctx)
+	u.cancelDownloadLock.Lock()
+	u.cancelDownload = cancel
+	u.cancelDownloadLock.Unlock()
+	defer func() {
+		u.cancelDownloadLock.Lock()
+		u.cancelDownload = nil
+		u.cancelDownloadLock.Unlock()
+		cancel()
+	}()
+
 	// Do a head first to check etag info
-	req, err := http.NewRequestWithContext(ctx, http.MethodHead, updateResp.UpdateURL, nil)
+	req, err := http.NewRequestWithContext(downloadCtx, http.MethodHead, updateResp.UpdateURL, nil)
 	if err != nil {
 		return err
 	}
 
 	// In case of slow downloads, continue the update check in the background
-	bgctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	bgctx, bgcancel := context.WithCancel(downloadCtx)
+	defer bgcancel()
 	go func() {
 		for {
 			select {
@@ -248,7 +261,20 @@ func cleanupOldDownloads(stageDir string) {
 }
 
 type Updater struct {
-	Store *store.Store
+	Store              *store.Store
+	cancelDownload     context.CancelFunc
+	cancelDownloadLock sync.Mutex
+}
+
+// CancelOngoingDownload cancels any currently running download
+func (u *Updater) CancelOngoingDownload() {
+	u.cancelDownloadLock.Lock()
+	defer u.cancelDownloadLock.Unlock()
+	if u.cancelDownload != nil {
+		slog.Info("cancelling ongoing update download")
+		u.cancelDownload()
+		u.cancelDownload = nil
+	}
 }
 
 func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) {