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 TestSubagentResumeBeforeHookRunsBeforeRestartedWorkerStarts(t *testing.T) {
82 root := t.TempDir()
83 started := make(chan *subagents.Agent, 2)
84 resumedTurnEnded := make(chan struct{})
85 tracker := subagents.NewCompletionTracker()
86 manager := subagents.New(subagents.Config{
87 Root: filepath.Join(root, "subagents"),
88 RepoRoot: root,
89 NewRunner: func(agent *subagents.Agent) subagents.Runner {
90 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
91 started <- agent
92 if agent.Resuming {
93 if agent.OnTurnEnd == nil {
94 return context.Canceled
95 }
96 agent.OnTurnEnd(1, "")
97 close(resumedTurnEnded)
98 return nil
99 }
100 <-ctx.Done()
101 return ctx.Err()
102 })
103 },
104 })
105 t.Cleanup(manager.StopAll)
106
107 first, err := manager.Spawn(context.Background(), "review the implementation")
108 if err != nil {
109 t.Fatal(err)
110 }
111 select {
112 case <-started:
113 case <-time.After(time.Second):
114 t.Fatal("initial subagent did not start")
115 }
116 if err := manager.Stop(first.ID); err != nil {
117 t.Fatal(err)
118 }
119 first.Wait()
120
121 const prompt = "continue the review"
122 var beforeAgent *subagents.Agent
123 var beforePrompt string
124 tool := &SubagentResumeTool{
125 Supervisor: manager,
126 Enabled: func() bool { return true },
127 BeforeResumed: func(agent *subagents.Agent, gotPrompt string, _ bool) func() {
128 beforeAgent = agent
129 beforePrompt = gotPrompt
130 return tracker.TrackFutureTurn(agent, gotPrompt, true)
131 },
132 OnResumed: func(*subagents.Agent, string, bool) {
133 t.Fatal("post-success callback should not run with a pre-resume hook")
134 },
135 }
136 args, err := json.Marshal(subagentResumeArgs{AgentID: first.ID, Prompt: prompt})
137 if err != nil {
138 t.Fatal(err)
139 }
140 if res, err := tool.Execute(context.Background(), args, nil); err != nil || res.IsError {
141 t.Fatalf("resume result = (%+v, %v), want success", res, err)
142 }
143 if beforeAgent == nil || beforeAgent == first || beforePrompt != prompt {
144 t.Fatalf("before hook = (%p, %q), want restarted worker and %q", beforeAgent, beforePrompt, prompt)
145 }
146 select {
147 case resumed := <-started:
148 if resumed.OnTurnEnd == nil {
149 t.Fatal("restarted worker started without its pre-registered callback")
150 }
151 case <-time.After(time.Second):
152 t.Fatal("resumed subagent did not start")
153 }
154 select {
155 case <-resumedTurnEnded:
156 case <-time.After(time.Second):
157 t.Fatal("resumed turn did not finish")
158 }
159 batch, err := tracker.WaitIdle(context.Background())
160 if err != nil || len(batch) != 1 || batch[0].Task != prompt {
161 t.Fatalf("resumed completion = (%+v, %v), want one %q completion", batch, err, prompt)
162 }
163}
164
165func TestSubagentResumeRestartsSessionWithFollowUp(t *testing.T) {
166 root := t.TempDir()
167 started := make(chan *subagents.Agent, 2)
168 manager := subagents.New(subagents.Config{
169 Root: filepath.Join(root, "subagents"),
170 RepoRoot: root,
171 NewRunner: func(agent *subagents.Agent) subagents.Runner {
172 return subagents.RunnerFunc(func(ctx context.Context, _ subagents.Sink) error {
173 started <- agent
174 <-ctx.Done()
175 return ctx.Err()
176 })
177 },
178 })
179 t.Cleanup(manager.StopAll)
180
181 first, err := manager.Spawn(context.Background(), "review the implementation")
182 if err != nil {
183 t.Fatal(err)
184 }
185 select {
186 case got := <-started:
187 if got != first {
188 t.Fatalf("first runner agent = %p, want %p", got, first)
189 }
190 case <-time.After(time.Second):
191 t.Fatal("initial subagent did not start")
192 }
193 if err := manager.Stop(first.ID); err != nil {
194 t.Fatal(err)
195 }
196 first.Wait()
197
198 const prompt = "I applied your review. What do you think now?"
199 var tracked *subagents.Agent
200 var trackedPrompt string
201 tool := &SubagentResumeTool{
202 Supervisor: manager,
203 Enabled: func() bool { return true },
204 OnResumed: func(agent *subagents.Agent, gotPrompt string, _ bool) {
205 tracked = agent
206 trackedPrompt = gotPrompt
207 },
208 }
209 args, err := json.Marshal(subagentResumeArgs{AgentID: first.ID, Prompt: prompt})
210 if err != nil {
211 t.Fatal(err)
212 }
213 res, err := tool.Execute(context.Background(), args, nil)
214 if err != nil {
215 t.Fatal(err)
216 }
217 if res.IsError {
218 t.Fatalf("unexpected tool error: %s", textResult(res.Content))
219 }
220 trackedID := ""
221 if tracked != nil {
222 trackedID = tracked.ID
223 }
224 if trackedID != first.ID || trackedPrompt != prompt {
225 t.Fatalf("resume callback = agent %q prompt %q, want agent %s prompt %q", trackedID, trackedPrompt, first.ID, prompt)
226 }
227
228 select {
229 case resumed := <-started:
230 if resumed.ID != first.ID {
231 t.Fatalf("resumed id = %q, want %q", resumed.ID, first.ID)
232 }
233 if resumed.SessionPath != first.SessionPath {
234 t.Fatalf("resumed session = %q, want %q", resumed.SessionPath, first.SessionPath)
235 }
236 resumePrompt, _ := resumed.ResumePromptInfo()
237 if !resumed.Resuming || resumePrompt != prompt {
238 t.Fatalf("resumed lifecycle = resuming %t prompt %q, want true and %q", resumed.Resuming, resumePrompt, prompt)
239 }
240 case <-time.After(time.Second):
241 t.Fatal("resumed subagent did not start")
242 }
243
244 var got struct {
245 Action string `json:"action"`
246 Agent struct {
247 ID string `json:"agent_id"`
248 } `json:"agent"`
249 }
250 if err := json.Unmarshal([]byte(textResult(res.Content)), &got); err != nil {
251 t.Fatalf("resume response JSON: %v", err)
252 }
253 if got.Action != "resumed" || got.Agent.ID != first.ID {
254 t.Fatalf("resume response = %+v, want resumed agent %s", got, first.ID)
255 }
256}