Fix task operation cancellation for lq0

8 files changed, 225 insertions, 24 deletions

diff --git a/internal/task/task.go b/internal/task/task.go
index ee9f0c9..e419f0a 100644
--- a/internal/task/task.go
+++ b/internal/task/task.go
@@ -63,7 +63,11 @@ type CompletionSources struct {
 }
 
 func run(args ...string) error {
-	_, err := RunArgs(context.Background(), args)
+	return runContext(context.Background(), args...)
+}
+
+func runContext(ctx context.Context, args ...string) error {
+	_, err := RunArgs(ctx, args)
 	return err
 }
 
@@ -342,6 +346,12 @@ func SetPriority(id int, priority string) error {
 
 // AddTags adds tags to the task with the given id.
 func AddTags(id int, tags []string) error {
+	return AddTagsContext(context.Background(), id, tags)
+}
+
+// AddTagsContext adds tags to the task with the given id using ctx for the
+// underlying Taskwarrior command.
+func AddTagsContext(ctx context.Context, id int, tags []string) error {
 	if id <= 0 {
 		return fmt.Errorf("invalid task ID: %d", id)
 	}
@@ -352,11 +362,17 @@ func AddTags(id int, tags []string) error {
 		}
 		args = append(args, t)
 	}
-	return run(args...)
+	return runContext(ctx, args...)
 }
 
 // RemoveTags removes tags from the task with the given id.
 func RemoveTags(id int, tags []string) error {
+	return RemoveTagsContext(context.Background(), id, tags)
+}
+
+// RemoveTagsContext removes tags from the task with the given id using ctx for
+// the underlying Taskwarrior command.
+func RemoveTagsContext(ctx context.Context, id int, tags []string) error {
 	if id <= 0 {
 		return fmt.Errorf("invalid task ID: %d", id)
 	}
@@ -367,7 +383,7 @@ func RemoveTags(id int, tags []string) error {
 		}
 		args = append(args, t)
 	}
-	return run(args...)
+	return runContext(ctx, args...)
 }
 
 // SetTags sets the tags of the task with the given id to exactly the provided set.
@@ -405,12 +421,12 @@ func SetTags(ctx context.Context, id int, tags []string) error {
 	}
 
 	if len(adds) > 0 {
-		if err := AddTags(id, adds); err != nil {
+		if err := AddTagsContext(ctx, id, adds); err != nil {
 			return err
 		}
 	}
 	if len(removes) > 0 {
-		if err := RemoveTags(id, removes); err != nil {
+		if err := RemoveTagsContext(ctx, id, removes); err != nil {
 			return err
 		}
 	}
@@ -439,10 +455,16 @@ func SetProject(id int, project string) error {
 
 // Annotate adds an annotation to the task with the given id.
 func Annotate(id int, text string) error {
+	return AnnotateContext(context.Background(), id, text)
+}
+
+// AnnotateContext adds an annotation to the task with the given id using ctx
+// for the underlying Taskwarrior command.
+func AnnotateContext(ctx context.Context, id int, text string) error {
 	if id <= 0 {
 		return fmt.Errorf("invalid task ID: %d", id)
 	}
-	return run(strconv.Itoa(id), "annotate", text)
+	return runContext(ctx, strconv.Itoa(id), "annotate", text)
 }
 
 // Denotate removes an annotation from the task with the given id.
@@ -450,6 +472,12 @@ func Annotate(id int, text string) error {
 // annotation text is matched exactly when provided. If text is empty, the
 // oldest annotation is removed.
 func Denotate(id int, text string) error {
+	return DenotateContext(context.Background(), id, text)
+}
+
+// DenotateContext removes an annotation from the task with the given id using
+// ctx for the underlying Taskwarrior command.
+func DenotateContext(ctx context.Context, id int, text string) error {
 	if id <= 0 {
 		return fmt.Errorf("invalid task ID: %d", id)
 	}
@@ -457,7 +485,7 @@ func Denotate(id int, text string) error {
 	if text != "" {
 		args = append(args, text)
 	}
-	return run(args...)
+	return runContext(ctx, args...)
 }
 
 // ReplaceAnnotations removes all existing annotations from the task with the
@@ -476,14 +504,14 @@ func ReplaceAnnotations(ctx context.Context, id int, text string) error {
 	}
 	anns := tasks[0].Annotations
 	for i := len(anns) - 1; i >= 0; i-- {
-		if err := Denotate(id, anns[i].Description); err != nil {
+		if err := DenotateContext(ctx, id, anns[i].Description); err != nil {
 			return err
 		}
 	}
 	if text == "" {
 		return nil
 	}
-	return Annotate(id, text)
+	return AnnotateContext(ctx, id, text)
 }
 
 // Edit opens the task in an editor for manual modification.
diff --git a/internal/task/task_test.go b/internal/task/task_test.go
index 59f94ba..2bca940 100644
--- a/internal/task/task_test.go
+++ b/internal/task/task_test.go
@@ -262,6 +262,134 @@ func TestExportReturnsCapturedErrorOutput(t *testing.T) {
 	}
 }
 
+func TestSetTagsHonorsContextDuringMutations(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	script := "#!/bin/sh\n" +
+		"for arg in \"$@\"; do\n" +
+		"  if [ \"$arg\" = modify ]; then\n" +
+		"    sleep 5\n" +
+		"    exit 0\n" +
+		"  fi\n" +
+		"done\n" +
+		"printf '%s\\n' '{\"id\":1,\"tags\":[\"old\"]}'\n"
+	if err := os.WriteFile(taskPath, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	origPath := os.Getenv("PATH")
+	t.Setenv("PATH", tmp+":"+origPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	err := SetTags(ctx, 1, []string{"new"})
+	elapsed := time.Since(start)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("SetTags error = %v, want context deadline exceeded", err)
+	}
+	if elapsed > time.Second {
+		t.Fatalf("SetTags took %s, expected prompt context cancellation", elapsed)
+	}
+}
+
+func TestSetTagsHonorsContextDuringRemovals(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	script := "#!/bin/sh\n" +
+		"for arg in \"$@\"; do\n" +
+		"  if [ \"$arg\" = modify ]; then\n" +
+		"    sleep 5\n" +
+		"    exit 0\n" +
+		"  fi\n" +
+		"done\n" +
+		"printf '%s\\n' '{\"id\":1,\"tags\":[\"old\"]}'\n"
+	if err := os.WriteFile(taskPath, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	origPath := os.Getenv("PATH")
+	t.Setenv("PATH", tmp+":"+origPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	err := SetTags(ctx, 1, nil)
+	elapsed := time.Since(start)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("SetTags error = %v, want context deadline exceeded", err)
+	}
+	if elapsed > time.Second {
+		t.Fatalf("SetTags took %s, expected prompt context cancellation", elapsed)
+	}
+}
+
+func TestReplaceAnnotationsHonorsContextDuringMutations(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	script := "#!/bin/sh\n" +
+		"for arg in \"$@\"; do\n" +
+		"  if [ \"$arg\" = denotate ] || [ \"$arg\" = annotate ]; then\n" +
+		"    sleep 5\n" +
+		"    exit 0\n" +
+		"  fi\n" +
+		"done\n" +
+		"printf '%s\\n' '{\"id\":1,\"annotations\":[{\"entry\":\"20260622T000000Z\",\"description\":\"old note\"}]}'\n"
+	if err := os.WriteFile(taskPath, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	origPath := os.Getenv("PATH")
+	t.Setenv("PATH", tmp+":"+origPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	err := ReplaceAnnotations(ctx, 1, "new note")
+	elapsed := time.Since(start)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("ReplaceAnnotations error = %v, want context deadline exceeded", err)
+	}
+	if elapsed > time.Second {
+		t.Fatalf("ReplaceAnnotations took %s, expected prompt context cancellation", elapsed)
+	}
+}
+
+func TestReplaceAnnotationsHonorsContextDuringAnnotate(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	script := "#!/bin/sh\n" +
+		"for arg in \"$@\"; do\n" +
+		"  if [ \"$arg\" = annotate ]; then\n" +
+		"    sleep 5\n" +
+		"    exit 0\n" +
+		"  fi\n" +
+		"done\n" +
+		"printf '%s\\n' '{\"id\":1,\"annotations\":[]}'\n"
+	if err := os.WriteFile(taskPath, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	origPath := os.Getenv("PATH")
+	t.Setenv("PATH", tmp+":"+origPath)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	err := ReplaceAnnotations(ctx, 1, "new note")
+	elapsed := time.Since(start)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("ReplaceAnnotations error = %v, want context deadline exceeded", err)
+	}
+	if elapsed > time.Second {
+		t.Fatalf("ReplaceAnnotations took %s, expected prompt context cancellation", elapsed)
+	}
+}
+
 func TestLoadCompletionSources(t *testing.T) {
 	tmp := t.TempDir()
 	taskPath := filepath.Join(tmp, "task")
diff --git a/internal/ui/handlers.go b/internal/ui/handlers.go
index 394aa80..45c0b10 100644
--- a/internal/ui/handlers.go
+++ b/internal/ui/handlers.go
@@ -47,14 +47,16 @@ func (m *Model) handleAnnotationMode(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
 		}
 
 		if m.replaceAnnotations {
-			ctx, cancel := taskExportContext()
+			ctx, cancel := m.taskExportContext()
 			defer cancel()
 			if err := task.ReplaceAnnotations(ctx, m.annotateID, value); err != nil {
 				return err
 			}
 			m.replaceAnnotations = false
 		} else {
-			if err := task.Annotate(m.annotateID, value); err != nil {
+			ctx, cancel := m.taskExportContext()
+			defer cancel()
+			if err := task.AnnotateContext(ctx, m.annotateID, value); err != nil {
 				return err
 			}
 		}
@@ -127,14 +129,18 @@ func (m *Model) handleTagsMode(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
 				}
 			}
 		}
-		if len(adds) > 0 {
-			if err := task.AddTags(m.tagsID, adds); err != nil {
-				return err
+		if len(adds) > 0 || len(removes) > 0 {
+			ctx, cancel := m.taskExportContext()
+			defer cancel()
+			if len(adds) > 0 {
+				if err := task.AddTagsContext(ctx, m.tagsID, adds); err != nil {
+					return err
+				}
 			}
-		}
-		if len(removes) > 0 {
-			if err := task.RemoveTags(m.tagsID, removes); err != nil {
-				return err
+			if len(removes) > 0 {
+				if err := task.RemoveTagsContext(ctx, m.tagsID, removes); err != nil {
+					return err
+				}
 			}
 		}
 		if err := m.reload(); err != nil {
diff --git a/internal/ui/keyactions.go b/internal/ui/keyactions.go
index 072ea20..c84c172 100644
--- a/internal/ui/keyactions.go
+++ b/internal/ui/keyactions.go
@@ -159,7 +159,7 @@ func (m *Model) handleUndo() (tea.Model, tea.Cmd) {
 			}
 			filters = append(filters, "status:"+restore.status)
 
-			ctx, cancel := taskExportContext()
+			ctx, cancel := m.taskExportContext()
 			tasks, err := task.Export(ctx, filters...)
 			cancel()
 			if err == nil && len(tasks) > 0 {
@@ -200,7 +200,7 @@ func (m *Model) deleteTaskWithUndo(tsk task.Task) (int, bool, error) {
 	recurring := isRecurringTask(tsk)
 	tasks := []task.Task{tsk}
 	if recurring {
-		ctx, cancel := taskExportContext()
+		ctx, cancel := m.taskExportContext()
 		series, err := task.RecurringSeries(ctx, recurringRootUUID(tsk))
 		cancel()
 		if err != nil {
@@ -495,7 +495,9 @@ func (m *Model) handleTagToProject() (tea.Model, tea.Cmd) {
 	}
 
 	// Remove the tag from the task
-	if err := task.RemoveTags(id, []string{firstTag}); err != nil {
+	ctx, cancel := m.taskExportContext()
+	defer cancel()
+	if err := task.RemoveTagsContext(ctx, id, []string{firstTag}); err != nil {
 		m.showError(err)
 		return m, nil
 	}
diff --git a/internal/ui/keyhandlers.go b/internal/ui/keyhandlers.go
index 0cfabec..16263db 100644
--- a/internal/ui/keyhandlers.go
+++ b/internal/ui/keyhandlers.go
@@ -184,6 +184,7 @@ func (m *Model) handleQuitKey() (tea.Model, tea.Cmd) {
 		m.reloadAndReport()
 		return m, nil
 	}
+	m.cancelTaskOperations()
 	return m, tea.Quit
 }
 
diff --git a/internal/ui/table.go b/internal/ui/table.go
index 9aeeda1..c948d89 100644
--- a/internal/ui/table.go
+++ b/internal/ui/table.go
@@ -228,6 +228,9 @@ type Model struct {
 	statusMsg string // temporary status message shown in status bar
 
 	helpViewport viewport.Model
+
+	taskContext       context.Context
+	cancelTaskContext context.CancelFunc
 }
 
 // editDoneMsg is emitted when the external editor process finishes.
@@ -262,8 +265,22 @@ type reloadData struct {
 	ultraFilterIDs []int
 }
 
-func taskExportContext() (context.Context, context.CancelFunc) {
-	return context.WithTimeout(context.Background(), taskExportTimeout)
+func (m *Model) initTaskContext() {
+	if m.taskContext != nil && m.cancelTaskContext != nil {
+		return
+	}
+	m.taskContext, m.cancelTaskContext = context.WithCancel(context.Background())
+}
+
+func (m *Model) cancelTaskOperations() {
+	if m.cancelTaskContext != nil {
+		m.cancelTaskContext()
+	}
+}
+
+func (m *Model) taskExportContext() (context.Context, context.CancelFunc) {
+	m.initTaskContext()
+	return context.WithTimeout(m.taskContext, taskExportTimeout)
 }
 
 // blinkInterval controls how quickly the row flashes when a task changes.
@@ -413,6 +430,7 @@ func (m *Model) startBlink(id int, markDone bool) tea.Cmd {
 // New creates a new UI model with the provided rows.
 func New(filters []string, browserCmd string) (Model, error) {
 	m := Model{filters: filters, browserCmd: browserCmd, agentFilterHotkey: "3", blinkState: blinkState{blinkEnabled: true}}
+	m.initTaskContext()
 	m.annotateInput = textinput.New()
 	m.annotateInput.Prompt = "annotation: "
 	m.descInput = textinput.New()
@@ -442,6 +460,7 @@ func New(filters []string, browserCmd string) (Model, error) {
 	m.theme = m.defaultTheme
 
 	if err := m.reload(); err != nil {
+		m.cancelTaskOperations()
 		return Model{}, err
 	}
 
@@ -491,7 +510,7 @@ func (m *Model) fetchTasks() (reloadData, error) {
 	// Always show only pending tasks by default.
 	filters := append([]string(nil), m.filters...)
 	filters = append(filters, "status:pending")
-	ctx, cancel := taskExportContext()
+	ctx, cancel := m.taskExportContext()
 	defer cancel()
 
 	tasks, err := task.Export(ctx, filters...)
diff --git a/internal/ui/table_test.go b/internal/ui/table_test.go
index c4f1e82..3a383e3 100644
--- a/internal/ui/table_test.go
+++ b/internal/ui/table_test.go
@@ -1,6 +1,8 @@
 package ui
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -1501,6 +1503,20 @@ func TestEscDoesNotQuitFromTable(t *testing.T) {
 	}
 }
 
+func TestQuitCancelsTaskExportContext(t *testing.T) {
+	m := Model{}
+	ctx, cancel := m.taskExportContext()
+	defer cancel()
+
+	_, cmd := m.handleQuitKey()
+	if cmd == nil {
+		t.Fatal("quit returned nil command; want tea.Quit")
+	}
+	if !errors.Is(ctx.Err(), context.Canceled) {
+		t.Fatalf("task export context error = %v, want context canceled", ctx.Err())
+	}
+}
+
 func TestEscDoesNotQuitUltraStartup(t *testing.T) {
 	tmp := t.TempDir()
 	taskPath := setupBasicTask(t, tmp)
diff --git a/internal/ui/ultra.go b/internal/ui/ultra.go
index 0d5f57b..9545c6f 100644
--- a/internal/ui/ultra.go
+++ b/internal/ui/ultra.go
@@ -1167,6 +1167,7 @@ func (m *Model) handleUltraExitKey(quit bool) (tea.Model, tea.Cmd) {
 	}
 	if m.ultraStartup {
 		if quit {
+			m.cancelTaskOperations()
 			return m, tea.Quit
 		}
 		return m, nil