From 976068369b3ec293bb2769975fcd64b4c50e1585 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 15 Nov 2023 16:59:49 -0800
Subject: [PATCH] stop all spinners on progress stop

---
 cmd/cmd.go           | 52 +++++++++++++++-----------------------------
 progress/progress.go |  6 +++++
 2 files changed, 24 insertions(+), 34 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index bc128d74..e3dd3837 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -53,7 +53,7 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 
 	bars := make(map[string]*progress.Bar)
 
-	status := "transferring context"
+	status := fmt.Sprintf("creating %s", args[0])
 	spinner := progress.NewSpinner(status)
 	p.Add(status, spinner)
 
@@ -72,6 +72,12 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	spinner.Stop()
+
+	status = "transferring context"
+	spinner = progress.NewSpinner(status)
+	p.Add(status, spinner)
+
 	for _, c := range commands {
 		switch c.Name {
 		case "model", "adapter":
@@ -133,10 +139,6 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if spinner != nil {
-		spinner.Stop()
-	}
-
 	return nil
 }
 
@@ -178,15 +180,13 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 
 	bars := make(map[string]*progress.Bar)
 
-	var status string
-	var spinner *progress.Spinner
+	status := fmt.Sprintf("pushing %s", args[0])
+	spinner := progress.NewSpinner(status)
+	p.Add(status, spinner)
 
 	fn := func(resp api.ProgressResponse) error {
 		if resp.Digest != "" {
-			if spinner != nil {
-				spinner.Stop()
-				spinner = nil
-			}
+			spinner.Stop()
 
 			bar, ok := bars[resp.Digest]
 			if !ok {
@@ -197,10 +197,7 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 
 			bar.Set(resp.Completed)
 		} else if status != resp.Status {
-			if spinner != nil {
-				spinner.Stop()
-				spinner = nil
-			}
+			spinner.Stop()
 
 			status = resp.Status
 			spinner = progress.NewSpinner(status)
@@ -215,10 +212,7 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if spinner != nil {
-		spinner.Stop()
-	}
-
+	spinner.Stop()
 	return nil
 }
 
@@ -379,15 +373,13 @@ func PullHandler(cmd *cobra.Command, args []string) error {
 
 	bars := make(map[string]*progress.Bar)
 
-	var status string
-	var spinner *progress.Spinner
+	status := fmt.Sprintf("pulling %s", args[0])
+	spinner := progress.NewSpinner(status)
+	p.Add(status, spinner)
 
 	fn := func(resp api.ProgressResponse) error {
 		if resp.Digest != "" {
-			if spinner != nil {
-				spinner.Stop()
-				spinner = nil
-			}
+			spinner.Stop()
 
 			bar, ok := bars[resp.Digest]
 			if !ok {
@@ -398,10 +390,7 @@ func PullHandler(cmd *cobra.Command, args []string) error {
 
 			bar.Set(resp.Completed)
 		} else if status != resp.Status {
-			if spinner != nil {
-				spinner.Stop()
-				spinner = nil
-			}
+			spinner.Stop()
 
 			status = resp.Status
 			spinner = progress.NewSpinner(status)
@@ -416,10 +405,6 @@ func PullHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if spinner != nil {
-		spinner.Stop()
-	}
-
 	return nil
 }
 
@@ -476,7 +461,6 @@ func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format st
 	defer p.Stop()
 
 	spinner := progress.NewSpinner("")
-	defer spinner.Stop()
 	p.Add("", spinner)
 
 	var latest api.GenerateResponse
diff --git a/progress/progress.go b/progress/progress.go
index e3d207cc..97acf82c 100644
--- a/progress/progress.go
+++ b/progress/progress.go
@@ -31,6 +31,12 @@ func NewProgress(w io.Writer) *Progress {
 }
 
 func (p *Progress) Stop() bool {
+	for _, state := range p.states {
+		if spinner, ok := state.(*Spinner); ok {
+			spinner.Stop()
+		}
+	}
+
 	if p.ticker != nil {
 		p.ticker.Stop()
 		p.ticker = nil