Fix annotation rollback cancellation (pq0)

2 files changed, 92 insertions, 10 deletions

diff --git a/internal/task/task.go b/internal/task/task.go
index ed0cd9c..604d4ca 100644
--- a/internal/task/task.go
+++ b/internal/task/task.go
@@ -616,20 +616,20 @@ func ReplaceAnnotations(ctx context.Context, id int, text string) error {
 	anns := tasks[0].Annotations
 	for i := len(anns) - 1; i >= 0; i-- {
 		if err := DenotateContext(ctx, id, anns[i].Description); err != nil {
-			return replaceAnnotationsError(ctx, id, anns, err)
+			return replaceAnnotationsError(id, anns, err)
 		}
 	}
 	if text == "" {
 		return nil
 	}
 	if err := AnnotateContext(ctx, id, text); err != nil {
-		return replaceAnnotationsError(ctx, id, anns, err)
+		return replaceAnnotationsError(id, anns, err)
 	}
 	return nil
 }
 
-func replaceAnnotationsError(ctx context.Context, id int, anns []Annotation, err error) error {
-	rollbackCtx, cancel := rollbackContext(ctx)
+func replaceAnnotationsError(id int, anns []Annotation, err error) error {
+	rollbackCtx, cancel := rollbackContext()
 	defer cancel()
 
 	if rollbackErr := restoreAnnotations(rollbackCtx, id, anns); rollbackErr != nil {
@@ -638,11 +638,8 @@ func replaceAnnotationsError(ctx context.Context, id int, anns []Annotation, err
 	return err
 }
 
-func rollbackContext(ctx context.Context) (context.Context, context.CancelFunc) {
-	if ctx.Err() != nil {
-		return ctx, func() {}
-	}
-	return context.WithTimeout(ctx, 5*time.Second)
+func rollbackContext() (context.Context, context.CancelFunc) {
+	return context.WithTimeout(context.Background(), 5*time.Second)
 }
 
 func restoreAnnotations(ctx context.Context, id int, anns []Annotation) error {
diff --git a/internal/task/task_test.go b/internal/task/task_test.go
index 3b3633e..556351f 100644
--- a/internal/task/task_test.go
+++ b/internal/task/task_test.go
@@ -465,10 +465,15 @@ exit 1
 func TestReplaceAnnotationsHonorsContextDuringMutations(t *testing.T) {
 	tmp := t.TempDir()
 	taskPath := filepath.Join(tmp, "task")
+	failPath := filepath.Join(tmp, "deadline-once")
 	script := "#!/bin/sh\n" +
+		"fail=" + shellQuote(failPath) + "\n" +
 		"for arg in \"$@\"; do\n" +
 		"  if [ \"$arg\" = denotate ] || [ \"$arg\" = annotate ]; then\n" +
-		"    sleep 5\n" +
+		"    if [ ! -e \"$fail\" ]; then\n" +
+		"      touch \"$fail\"\n" +
+		"      sleep 5\n" +
+		"    fi\n" +
 		"    exit 0\n" +
 		"  fi\n" +
 		"done\n" +
@@ -526,6 +531,86 @@ func TestReplaceAnnotationsHonorsContextDuringAnnotate(t *testing.T) {
 	}
 }
 
+func TestReplaceAnnotationsRestoresSnapshotAfterDenotateDeadline(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	statePath := filepath.Join(tmp, "annotations.txt")
+	failPath := filepath.Join(tmp, "deadline-once")
+	if err := os.WriteFile(statePath, []byte("first note\nsecond note\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	script := fakeAnnotationTaskScript(statePath, `
+if [ "$2" = denotate ] && [ "$3" = "first note" ] && [ ! -e `+shellQuote(failPath)+` ]; then
+  touch `+shellQuote(failPath)+`
+  sleep 5
+  exit 0
+fi
+`)
+	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, "replacement 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 plus rollback", elapsed)
+	}
+	if got := readLinesFile(t, statePath); strings.Join(got, "|") != "first note|second note" {
+		t.Fatalf("annotations after rollback = %#v, want original annotations", got)
+	}
+}
+
+func TestReplaceAnnotationsRestoresSnapshotAfterAnnotateDeadline(t *testing.T) {
+	tmp := t.TempDir()
+	taskPath := filepath.Join(tmp, "task")
+	statePath := filepath.Join(tmp, "annotations.txt")
+	failPath := filepath.Join(tmp, "deadline-once")
+	if err := os.WriteFile(statePath, []byte("first note\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	script := fakeAnnotationTaskScript(statePath, `
+if [ "$2" = annotate ] && [ "$3" = "replacement note" ] && [ ! -e `+shellQuote(failPath)+` ]; then
+  touch `+shellQuote(failPath)+`
+  sleep 5
+  exit 0
+fi
+`)
+	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, "replacement 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 plus rollback", elapsed)
+	}
+	if got := readLinesFile(t, statePath); strings.Join(got, "|") != "first note" {
+		t.Fatalf("annotations after rollback = %#v, want original annotations", got)
+	}
+}
+
 func TestReplaceAnnotationsRestoresSnapshotAfterDenotateFailure(t *testing.T) {
 	tmp := t.TempDir()
 	taskPath := filepath.Join(tmp, "task")