1package tools
2
3import (
4 "context"
5 "encoding/json"
6 "path/filepath"
7 "testing"
8 "time"
9
10 "github.com/bnema/zut/packages/agent/subagents"
11)
12
13func TestSubagentStopTerminatesStuckWorker(t *testing.T) {
14 root := t.TempDir()
15 started := make(chan struct{}, 1)
16 manager := subagents.New(subagents.Config{
17 Root: filepath.Join(root, "subagents"),
18 RepoRoot: root,
19 NewRunner: func(*subagents.Agent) subagents.Runner {
20 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
21 started <- struct{}{}
22 <-ctx.Done()
23 return ctx.Err()
24 })
25 },
26 })
27 t.Cleanup(manager.StopAll)
28
29 agent, err := manager.Spawn(context.Background(), "investigate a stuck operation")
30 if err != nil {
31 t.Fatal(err)
32 }
33 select {
34 case <-started:
35 case <-time.After(time.Second):
36 t.Fatal("subagent did not start")
37 }
38
39 var tracked *subagents.Agent
40 tool := &SubagentStopTool{
41 Supervisor: manager,
42 Enabled: func() bool { return true },
43 OnStopRequested: func(got *subagents.Agent) {
44 tracked = got
45 },
46 }
47 args, err := json.Marshal(subagentStopArgs{AgentID: agent.ID})
48 if err != nil {
49 t.Fatal(err)
50 }
51 res, err := tool.Execute(context.Background(), args, nil)
52 if err != nil {
53 t.Fatal(err)
54 }
55 if res.IsError {
56 t.Fatalf("unexpected tool error: %s", textResult(res.Content))
57 }
58 if tracked != agent {
59 t.Fatalf("stop callback agent = %p, want %p", tracked, agent)
60 }
61 agent.Wait()
62 if got := agent.Status(); got != subagents.StatusKilled {
63 t.Fatalf("agent status = %s, want %s", got, subagents.StatusKilled)
64 }
65
66 var got struct {
67 Action string `json:"action"`
68 Agent struct {
69 ID string `json:"agent_id"`
70 State string `json:"state"`
71 } `json:"agent"`
72 }
73 if err := json.Unmarshal([]byte(textResult(res.Content)), &got); err != nil {
74 t.Fatalf("stop response JSON: %v", err)
75 }
76 if got.Action != "stop_requested" || got.Agent.ID != agent.ID || got.Agent.State != "cancelled" {
77 t.Fatalf("stop response = %+v, want stop_requested cancelled agent %s", got, agent.ID)
78 }
79}
80
81func TestSubagentResumeRestartsSessionWithFollowUp(t *testing.T) {
82 root := t.TempDir()
83 started := make(chan *subagents.Agent, 2)
84 manager := subagents.New(subagents.Config{
85 Root: filepath.Join(root, "subagents"),
86 RepoRoot: root,
87 NewRunner: func(agent *subagents.Agent) subagents.Runner {
88 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
89 started <- agent
90 <-ctx.Done()
91 return ctx.Err()
92 })
93 },
94 })
95 t.Cleanup(manager.StopAll)
96
97 first, err := manager.Spawn(context.Background(), "review the implementation")
98 if err != nil {
99 t.Fatal(err)
100 }
101 select {
102 case got := <-started:
103 if got != first {
104 t.Fatalf("first runner agent = %p, want %p", got, first)
105 }
106 case <-time.After(time.Second):
107 t.Fatal("initial subagent did not start")
108 }
109 if err := manager.Stop(first.ID); err != nil {
110 t.Fatal(err)
111 }
112 first.Wait()
113
114 const prompt = "I applied your review. What do you think now?"
115 var tracked *subagents.Agent
116 var trackedPrompt string
117 tool := &SubagentResumeTool{
118 Supervisor: manager,
119 Enabled: func() bool { return true },
120 OnResumed: func(agent *subagents.Agent, gotPrompt string) {
121 tracked = agent
122 trackedPrompt = gotPrompt
123 },
124 }
125 args, err := json.Marshal(subagentResumeArgs{AgentID: first.ID, Prompt: prompt})
126 if err != nil {
127 t.Fatal(err)
128 }
129 res, err := tool.Execute(context.Background(), args, nil)
130 if err != nil {
131 t.Fatal(err)
132 }
133 if res.IsError {
134 t.Fatalf("unexpected tool error: %s", textResult(res.Content))
135 }
136 trackedID := ""
137 if tracked != nil {
138 trackedID = tracked.ID
139 }
140 if trackedID != first.ID || trackedPrompt != prompt {
141 t.Fatalf("resume callback = agent %q prompt %q, want agent %s prompt %q", trackedID, trackedPrompt, first.ID, prompt)
142 }
143
144 select {
145 case resumed := <-started:
146 if resumed.ID != first.ID {
147 t.Fatalf("resumed id = %q, want %q", resumed.ID, first.ID)
148 }
149 if resumed.SessionPath != first.SessionPath {
150 t.Fatalf("resumed session = %q, want %q", resumed.SessionPath, first.SessionPath)
151 }
152 resumePrompt, _ := resumed.ResumePromptInfo()
153 if !resumed.Resuming || resumePrompt != prompt {
154 t.Fatalf("resumed lifecycle = resuming %t prompt %q, want true and %q", resumed.Resuming, resumePrompt, prompt)
155 }
156 case <-time.After(time.Second):
157 t.Fatal("resumed subagent did not start")
158 }
159
160 var got struct {
161 Action string `json:"action"`
162 Agent struct {
163 ID string `json:"agent_id"`
164 } `json:"agent"`
165 }
166 if err := json.Unmarshal([]byte(textResult(res.Content)), &got); err != nil {
167 t.Fatalf("resume response JSON: %v", err)
168 }
169 if got.Action != "resumed" || got.Agent.ID != first.ID {
170 t.Fatalf("resume response = %+v, want resumed agent %s", got, first.ID)
171 }
172}