From ed2e06e75c3f4aafc8acbf91d9cb13b8f04eda83 Mon Sep 17 00:00:00 2001
From: Edmund Wang <284339473+wangedmund77-cmyk@users.noreply.github.com>
Date: Sat, 11 Jul 2026 09:13:25 +0800
Subject: [PATCH] fix: respect cancellation after task action returns

---
 workspace/task/runner.go      |  5 ++++-
 workspace/task/runner_test.go | 22 +++++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/workspace/task/runner.go b/workspace/task/runner.go
index 0ba04c6..d2531bc 100644
--- a/workspace/task/runner.go
+++ b/workspace/task/runner.go
@@ -51,6 +51,9 @@ func (r *Runner) Run(ctx context.Context) {
 			task.mu.Unlock()
 
 			err := task.Action(ctx)
+			if err == nil && ctx.Err() != nil {
+				err = ctx.Err()
+			}
 
 			task.mu.Lock()
 			task.FinishedAt = time.Now()
@@ -66,4 +69,4 @@ func (r *Runner) Run(ctx context.Context) {
 		}(t)
 	}
 	wg.Wait()
-}
\ No newline at end of file
+}
diff --git a/workspace/task/runner_test.go b/workspace/task/runner_test.go
index d02f599..658687e 100644
--- a/workspace/task/runner_test.go
+++ b/workspace/task/runner_test.go
@@ -98,4 +98,24 @@ func TestRunner_Run_Cancel(t *testing.T) {
 	if err := t1.GetErr(); err != context.Canceled {
 		t.Errorf("expected t1 error to be context.Canceled, got %v", err)
 	}
-}
\ No newline at end of file
+}
+
+func TestRunner_Run_CancelDuringAction(t *testing.T) {
+	runner := NewRunner()
+	ctx, cancel := context.WithCancel(context.Background())
+
+	task := NewTask("cancel-during-action", func(ctx context.Context) error {
+		cancel()
+		return nil
+	})
+
+	runner.AddTask(task)
+	runner.Run(ctx)
+
+	if state := task.GetState(); state != StateFailed {
+		t.Errorf("expected task to fail after context cancellation, got %s", state)
+	}
+	if err := task.GetErr(); err != context.Canceled {
+		t.Errorf("expected task error to be context.Canceled, got %v", err)
+	}
+}