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>
This commit is contained in:
dj-oyu 2026-03-01 05:14:51 +09:00
parent fb9e8a3618
commit b7ed441ef4
2 changed files with 18 additions and 7 deletions

View file

@ -28,6 +28,7 @@ type SubagentTask struct {
Iterations int `json:"-"` Iterations int `json:"-"`
ToolCalls int `json:"-"` ToolCalls int `json:"-"`
ToolStats map[string]int `json:"-"` ToolStats map[string]int `json:"-"`
cancel context.CancelFunc
} }
type SubagentManager struct { type SubagentManager struct {
@ -126,10 +127,13 @@ func (sm *SubagentManager) Spawn(
// Start task in background with a detached context. // Start task in background with a detached context.
// The spawned goroutine must outlive the parent (e.g. heartbeat session) // The spawned goroutine must outlive the parent (e.g. heartbeat session)
// which may finish before the subagent completes. // 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) sm.wg.Add(1)
go func() { go func() {
defer sm.wg.Done() defer sm.wg.Done()
sm.runTask(context.Background(), subagentTask, preset, callback) sm.runTask(spawnCtx, subagentTask, preset, callback)
}() }()
if label != "" { if label != "" {
@ -379,6 +383,16 @@ func (sm *SubagentManager) WaitAll() {
sm.wg.Wait() 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) { func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) {
sm.mu.RLock() sm.mu.RLock()
defer sm.mu.RUnlock() defer sm.mu.RUnlock()

View file

@ -196,10 +196,7 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) {
bp := newBlockingProvider() bp := newBlockingProvider()
mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{})
ctx, cancel := context.WithCancel(context.Background()) _, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil)
defer cancel()
_, err := mgr.Spawn(ctx, "long task", "cancel-me", "", "cli", "direct", "", nil)
if err != nil { if err != nil {
t.Fatalf("Spawn() error: %v", err) 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") t.Fatal("timed out waiting for blockingProvider to enter Chat")
} }
// Now cancel — the LLM call unblocks with ctx.Err(). // Cancel via CancelTask — the spawned goroutine's detached context is canceled.
cancel() mgr.CancelTask("subagent-1")
// Collect events until agent_gc. // Collect events until agent_gc.
var events []orch.Event var events []orch.Event