1package swarm
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "os"
8 "path/filepath"
9 "strings"
10 "testing"
11 "time"
12)
13
14// TestSwarmAgentArgs locks in the exact flag set the subprocess
15// runner uses to start a swarm agent in daemon mode. Past
16// regressions in this area:
17//
18// - "--no-sess" instead of "--no-session" (old print-mode runner):
19// every spawned agent died with "unknown flag" before it could
20// talk to the model.
21//
22// - Forgetting --cwd: the child resolved tools against the parent
23// zot's working directory, defeating the whole point of the
24// worktree isolation.
25//
26// - Forgetting --session: a daemon-mode agent without a session
27// file would lose context between follow-up turns, making
28// "send another message" mostly useless.
29//
30// The test asserts the load-bearing pieces are present in plausible
31// positions. If a flag is renamed, update both the runner and this
32// test so we notice immediately.
33func TestCredentialStdinHelperProcess(t *testing.T) {
34 if os.Getenv("ZOT_SWARM_CREDENTIAL_HELPER") != "1" {
35 return
36 }
37 if os.Getenv("ZOT_SWARM_CREDENTIAL_STDIN") != "1" {
38 os.Exit(2)
39 }
40 var credential Credential
41 if err := json.NewDecoder(os.Stdin).Decode(&credential); err != nil {
42 os.Exit(3)
43 }
44 if credential.Value != "inherited-secret" || credential.Method != "apikey" {
45 os.Exit(4)
46 }
47 for _, arg := range os.Args {
48 if strings.Contains(arg, credential.Value) {
49 os.Exit(5)
50 }
51 }
52 for _, env := range os.Environ() {
53 if strings.Contains(env, credential.Value) {
54 os.Exit(6)
55 }
56 }
57 fmt.Println(`{"type":"agent_stopped","data":{"reason":"completed"}}`)
58 os.Exit(0)
59}
60
61func TestExecRunnerTransfersCredentialOnlyOnStdin(t *testing.T) {
62 t.Setenv("ZOT_SWARM_CREDENTIAL_HELPER", "1")
63 root := t.TempDir()
64 a := &Agent{
65 ID: "credential-test",
66 Dir: root,
67 Provider: "custom",
68 SessionPath: filepath.Join(root, "session.json"),
69 InboxPath: filepath.Join(root, "inbox"),
70 EventLogPath: filepath.Join(root, "events.jsonl"),
71 }
72 calls := 0
73 r := &execRunner{
74 agent: a,
75 Command: []string{os.Args[0], "-test.run=^TestCredentialStdinHelperProcess$"},
76 resolveCredential: func(context.Context, string) (Credential, error) {
77 calls++
78 return Credential{Value: "inherited-secret", Method: "apikey"}, nil
79 },
80 }
81 if err := r.Run(context.Background(), agentSink{a: a}); err != nil {
82 t.Fatal(err)
83 }
84 if calls != 1 {
85 t.Fatalf("credential resolver called %d times, want 1", calls)
86 }
87}
88
89func TestApplyEventPreservesMultilineRolesAndFinalAssistant(t *testing.T) {
90 a := &Agent{}
91 sink := agentSink{a: a}
92 applyEventToSink(NewEvent("user_message", map[string]any{"content": []any{
93 map[string]any{"type": "text", "text": "first line\n\nlast line"},
94 }}), sink)
95 applyEventToSink(NewEvent("assistant_message", map[string]any{"content": []any{
96 map[string]any{"type": "text", "text": "complete\nanswer"},
97 }}), sink)
98
99 snap := a.Snapshot()
100 wantLines := []string{"user: first line", "user: ", "user: last line", "complete", "answer"}
101 if strings.Join(snap.Lines, "\n") != strings.Join(wantLines, "\n") {
102 t.Fatalf("transcript = %#v, want %#v", snap.Lines, wantLines)
103 }
104 if snap.LastAssistant != "complete\nanswer" {
105 t.Fatalf("last assistant = %q", snap.LastAssistant)
106 }
107}
108
109func TestSwarmAgentArgs(t *testing.T) {
110 args := swarmAgentArgs(swarmAgentArgsOpts{
111 Exe: "/path/to/zot",
112 Dir: "/tmp/worktree",
113 SessionPath: "/tmp/state/session.json",
114 InboxPath: "/tmp/state/in.sock",
115 Task: "do the thing",
116 })
117 if len(args) < 7 {
118 t.Fatalf("argv unexpectedly short: %v", args)
119 }
120 if args[0] != "/path/to/zot" {
121 t.Fatalf("argv[0] = %q; want the binary path", args[0])
122 }
123 // The task must come last so anything that looks flag-like in
124 // the task body doesn't get interpreted as a flag.
125 if args[len(args)-1] != "do the thing" {
126 t.Fatalf("task should be last positional; got %v", args)
127 }
128
129 mustHave := map[string]string{
130 "--swarm-agent": "/tmp/state/in.sock",
131 "--session": "/tmp/state/session.json",
132 "--cwd": "/tmp/worktree",
133 }
134 for flag, value := range mustHave {
135 i := indexOf(args, flag)
136 if i < 0 {
137 t.Errorf("argv missing %q: %v", flag, args)
138 continue
139 }
140 if i+1 >= len(args) || args[i+1] != value {
141 t.Errorf("argv %q value = %q; want %q", flag, safeAt(args, i+1), value)
142 }
143 }
144
145 // Reject prior bad flags explicitly so a future revert is caught.
146 joined := strings.Join(args, " ")
147 for _, bad := range []string{"--print", "--no-sess ", "--no-session"} {
148 if strings.Contains(joined, bad) {
149 t.Fatalf("argv contains stale/wrong flag %q: %s", bad, joined)
150 }
151 }
152}
153
154// TestSwarmAgentArgsEmptyTaskOmitsPositional makes sure that when the
155// agent is being adopted (no fresh task) we don't pass an empty
156// positional which the arg parser would treat as a real prompt.
157func TestSwarmAgentArgsEmptyTaskOmitsPositional(t *testing.T) {
158 args := swarmAgentArgs(swarmAgentArgsOpts{
159 Exe: "/zot", Dir: "/wt", SessionPath: "/s.json", InboxPath: "/in.sock",
160 })
161 for _, a := range args {
162 if a == "" {
163 t.Fatalf("argv contains an empty positional: %v", args)
164 }
165 }
166 // last arg should be a real flag value, not a stray positional
167 if a := args[len(args)-1]; strings.HasPrefix(a, "--") {
168 t.Fatalf("argv ends on a flag with no value: %v", args)
169 }
170}
171
172// TestDefaultChildArgsSpawnIncludesTask pins the spawn shape: a
173// fresh (non-resuming) Agent produces argv that ends with the
174// original task as a positional, so the child runs it as the
175// initial user turn.
176func TestDefaultChildArgsSpawnIncludesTask(t *testing.T) {
177 a := &Agent{Dir: "/wt", Task: "do thing"}
178 args := defaultChildArgs("/zot", a, "/s.json", "/in.sock")
179 if got := args[len(args)-1]; got != "do thing" {
180 t.Fatalf("spawn argv last = %q; want %q\n%v", got, "do thing", args)
181 }
182}
183
184// TestDefaultChildArgsResumeOmitsTask is the regression for the
185// "agent busy; send 'cancel' first" error: when an Agent is being
186// resumed (Resuming==true), the child argv MUST NOT include the
187// original Task as a positional. Otherwise the child fires the task
188// as a fresh user turn on every resume, racing with whatever the
189// user types next via the inbox.
190func TestDefaultChildArgsResumeOmitsTask(t *testing.T) {
191 a := &Agent{Dir: "/wt", Task: "do thing", Resuming: true}
192 args := defaultChildArgs("/zot", a, "/s.json", "/in.sock")
193 for _, v := range args {
194 if v == "do thing" {
195 t.Fatalf("resume argv contains the task; it would re-fire as a duplicate turn\n%v", args)
196 }
197 }
198 // And no trailing positional at all: the last arg should be a
199 // flag value, not a stray empty string.
200 if got := args[len(args)-1]; got == "" {
201 t.Fatalf("resume argv ends with an empty positional: %v", args)
202 }
203 if strings.HasPrefix(args[len(args)-1], "--") {
204 t.Fatalf("resume argv ends on a bare flag (no value): %v", args)
205 }
206}
207
208func TestNotifyPromptTurnEndIgnoresProviderToolLoopTurnEnd(t *testing.T) {
209 called := make(chan struct{}, 1)
210 a := &Agent{}
211 a.SetOnTurnEnd(func(step int, errMsg string) {
212 called <- struct{}{}
213 })
214
215 notifyPromptTurnEnd(a, NewEvent("turn_end", map[string]any{"stop": "tool_use"}))
216
217 select {
218 case <-called:
219 t.Fatal("OnTurnEnd fired for provider/tool-loop turn_end without step")
220 case <-time.After(50 * time.Millisecond):
221 }
222}
223
224func TestNotifyPromptTurnEndFiresForDaemonPromptCompletion(t *testing.T) {
225 type got struct {
226 step int
227 err string
228 }
229 called := make(chan got, 1)
230 a := &Agent{}
231 a.SetOnTurnEnd(func(step int, errMsg string) {
232 called <- got{step: step, err: errMsg}
233 })
234
235 notifyPromptTurnEnd(a, NewEvent("turn_end", map[string]any{"step": float64(2), "error": "boom"}))
236
237 select {
238 case g := <-called:
239 if g.step != 2 || g.err != "boom" {
240 t.Fatalf("callback = (%d, %q); want (2, boom)", g.step, g.err)
241 }
242 case <-time.After(time.Second):
243 t.Fatal("OnTurnEnd did not fire for prompt-level turn_end with step")
244 }
245}
246
247func indexOf(xs []string, want string) int {
248 for i, x := range xs {
249 if x == want {
250 return i
251 }
252 }
253 return -1
254}
255
256func safeAt(xs []string, i int) string {
257 if i < 0 || i >= len(xs) {
258 return ""
259 }
260 return xs[i]
261}