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