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:
parent
8e97860c82
commit
354497f8a8
2 changed files with 18 additions and 7 deletions
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue