From 354497f8a8c1e901b388d77eeb4e26b221f31ee8 Mon Sep 17 00:00:00 2001
From: dj-oyu <68707227+dj-oyu@users.noreply.github.com>
Date: Sun, 1 Mar 2026 05:14:51 +0900
Subject: [PATCH] fix: update cancel test for detached spawn context

Spawn now uses a detached context. Add CancelTask() method with per-task
cancel func so subagents can still be explicitly stopped. Update
TestSubagentManager_Spawn_CancelledDuringExecution to use CancelTask()
instead of parent context cancellation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 pkg/tools/subagent.go               | 16 +++++++++++++++-
 pkg/tools/subagent_reporter_test.go |  9 +++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go
index f4acb4a4f..79b6300e2 100644
--- a/pkg/tools/subagent.go
+++ b/pkg/tools/subagent.go
@@ -28,6 +28,7 @@ type SubagentTask struct {
 	Iterations    int            `json:"-"`
 	ToolCalls     int            `json:"-"`
 	ToolStats     map[string]int `json:"-"`
+	cancel        context.CancelFunc
 }
 
 type SubagentManager struct {
@@ -126,10 +127,13 @@ func (sm *SubagentManager) Spawn(
 	// Start task in background with a detached context.
 	// The spawned goroutine must outlive the parent (e.g. heartbeat session)
 	// which may finish before the subagent completes.
+	// The cancel func is stored on the task so CancelTask() can stop it.
+	spawnCtx, spawnCancel := context.WithCancel(context.Background())
+	subagentTask.cancel = spawnCancel
 	sm.wg.Add(1)
 	go func() {
 		defer sm.wg.Done()
-		sm.runTask(context.Background(), subagentTask, preset, callback)
+		sm.runTask(spawnCtx, subagentTask, preset, callback)
 	}()
 
 	if label != "" {
@@ -379,6 +383,16 @@ func (sm *SubagentManager) WaitAll() {
 	sm.wg.Wait()
 }
 
+// CancelTask cancels the context for a running subagent task.
+func (sm *SubagentManager) CancelTask(taskID string) {
+	sm.mu.RLock()
+	task, ok := sm.tasks[taskID]
+	sm.mu.RUnlock()
+	if ok && task.cancel != nil {
+		task.cancel()
+	}
+}
+
 func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) {
 	sm.mu.RLock()
 	defer sm.mu.RUnlock()
diff --git a/pkg/tools/subagent_reporter_test.go b/pkg/tools/subagent_reporter_test.go
index f378b3060..0dcc2d6ac 100644
--- a/pkg/tools/subagent_reporter_test.go
+++ b/pkg/tools/subagent_reporter_test.go
@@ -196,10 +196,7 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) {
 	bp := newBlockingProvider()
 	mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{})
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	_, err := mgr.Spawn(ctx, "long task", "cancel-me", "", "cli", "direct", "", nil)
+	_, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil)
 	if err != nil {
 		t.Fatalf("Spawn() error: %v", err)
 	}
@@ -211,8 +208,8 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) {
 		t.Fatal("timed out waiting for blockingProvider to enter Chat")
 	}
 
-	// Now cancel — the LLM call unblocks with ctx.Err().
-	cancel()
+	// Cancel via CancelTask — the spawned goroutine's detached context is canceled.
+	mgr.CancelTask("subagent-1")
 
 	// Collect events until agent_gc.
 	var events []orch.Event