1package agent
2
3import (
4 "archive/tar"
5 "bytes"
6 "compress/gzip"
7 "errors"
8 "io"
9 "net/http"
10 "net/http/httptest"
11 "net/url"
12 "os"
13 "path/filepath"
14 "strings"
15 "testing"
16
17 "github.com/patriceckhart/zot/packages/agent/tools"
18 "github.com/patriceckhart/zot/packages/tui"
19)
20
21func TestReadZotfileConsentKeyAcceptsOnlyYesNo(t *testing.T) {
22 tests := []struct {
23 name string
24 input []byte
25 want byte
26 }{
27 {name: "lowercase yes", input: []byte{'y', '\r'}, want: 'y'},
28 {name: "uppercase yes", input: []byte{'Y', '\n'}, want: 'y'},
29 {name: "lowercase no", input: []byte{'n', '\r'}, want: 'n'},
30 {name: "uppercase no", input: []byte{'N', '\n'}, want: 'n'},
31 {name: "ignore unrelated keys", input: append([]byte("x\r\n\x1b[A"), 'y', '\r'), want: 'y'},
32 {name: "keep first valid answer", input: []byte{'n', 'y', '\r'}, want: 'n'},
33 {name: "backspace changes answer", input: []byte{'y', 0x7f, 'n', '\r'}, want: 'n'},
34 {name: "ctrl h changes answer", input: []byte{'n', '\b', 'y', '\r'}, want: 'y'},
35 {name: "backspace without answer is ignored", input: []byte{0x7f, 'y', '\r'}, want: 'y'},
36 }
37 for _, tt := range tests {
38 t.Run(tt.name, func(t *testing.T) {
39 var out bytes.Buffer
40 got, err := readZotfileConsentKey(bytes.NewReader(tt.input), func(answer byte) {
41 if answer == 0 {
42 out.Truncate(out.Len() - 1)
43 return
44 }
45 out.WriteByte(answer)
46 })
47 if err != nil {
48 t.Fatal(err)
49 }
50 if got != tt.want {
51 t.Fatalf("answer = %q, want %q", got, tt.want)
52 }
53 if out.String() != string(tt.want) {
54 t.Fatalf("echo = %q, want %q", out.String(), tt.want)
55 }
56 })
57 }
58}
59
60func TestReadZotfileConsentKeyWaitsForEnter(t *testing.T) {
61 var out bytes.Buffer
62 if _, err := readZotfileConsentKey(strings.NewReader("y"), func(answer byte) {
63 out.WriteByte(answer)
64 }); !errors.Is(err, io.EOF) {
65 t.Fatalf("error = %v, want EOF before Enter", err)
66 }
67 if out.String() != "y" {
68 t.Fatalf("echo = %q, want y", out.String())
69 }
70}
71
72func TestReadZotfileConsentKeyInterruptsOnCtrlC(t *testing.T) {
73 if _, err := readZotfileConsentKey(bytes.NewReader([]byte{'x', 3, 'y'}), nil); !errors.Is(err, errZotfileConsentInterrupted) {
74 t.Fatalf("error = %v, want interrupted", err)
75 }
76}
77
78func TestReadZotfileConsentKeyReportsInputError(t *testing.T) {
79 if _, err := readZotfileConsentKey(strings.NewReader("xxx"), nil); !errors.Is(err, io.EOF) {
80 t.Fatalf("error = %v, want EOF", err)
81 }
82}
83
84func writeTestZotfile(t *testing.T, manifest string) string {
85 t.Helper()
86 dir := t.TempDir()
87 if err := os.WriteFile(filepath.Join(dir, "manifest.json"), []byte(manifest), 0o600); err != nil {
88 t.Fatal(err)
89 }
90 if err := os.WriteFile(filepath.Join(dir, "AGENT.md"), []byte("Be useful."), 0o600); err != nil {
91 t.Fatal(err)
92 }
93 return dir
94}
95
96func TestApplyZotfileAgentPromptPreservesEmptyReplacement(t *testing.T) {
97 args := Args{SystemPrompt: "cli prompt"}
98 applyZotfileAgentPrompt(&args, true, " \n")
99 if !args.SystemPromptSet {
100 t.Fatal("SystemPromptSet = false; want true for an empty replacement AGENT.md")
101 }
102 if args.SystemPrompt != "" {
103 t.Fatalf("SystemPrompt = %q; want empty", args.SystemPrompt)
104 }
105}
106
107func TestLoadZotfileAllowsMissingAgentMD(t *testing.T) {
108 dir := t.TempDir()
109 if err := os.WriteFile(filepath.Join(dir, "manifest.json"), []byte(`{"zotfile":1,"name":"test"}`), 0o600); err != nil {
110 t.Fatal(err)
111 }
112 zf, cleanup, err := loadZotfile(dir)
113 if cleanup != nil {
114 defer cleanup()
115 }
116 if err != nil {
117 t.Fatalf("load without AGENT.md: %v", err)
118 }
119 if zf.Manifest.Name != "test" {
120 t.Fatalf("name = %q, want test", zf.Manifest.Name)
121 }
122}
123
124func TestValidateZotfileDirRejectsAgentMDDirectory(t *testing.T) {
125 dir := t.TempDir()
126 if err := os.Mkdir(filepath.Join(dir, "AGENT.md"), 0o700); err != nil {
127 t.Fatal(err)
128 }
129 if err := validateZotfileDir(dir); err == nil || !strings.Contains(err.Error(), "regular file") {
130 t.Fatalf("unexpected error: %v", err)
131 }
132}
133
134func TestLoadZotfileRejectsUnenforcedPermissions(t *testing.T) {
135 for _, field := range []string{
136 `"net":{"allow":["example.com"]}`,
137 `"env":{"read":["HOME"]}`,
138 } {
139 dir := writeTestZotfile(t, `{"zotfile":1,"name":"test","permissions":{`+field+`}}`)
140 if _, _, err := loadZotfile(dir); err == nil {
141 t.Fatalf("manifest with %s was accepted", field)
142 }
143 }
144}
145
146func TestLoadZotfileBashModes(t *testing.T) {
147 for _, mode := range []string{"none", "ask", "allowlist"} {
148 manifest := `{"zotfile":1,"name":"test","permissions":{"bash":{"mode":"` + mode + `"}}}`
149 if mode == "allowlist" {
150 manifest = `{"zotfile":1,"name":"test","permissions":{"bash":{"mode":"allowlist","allow":["git"]}}}`
151 }
152 dir := writeTestZotfile(t, manifest)
153 if _, _, err := loadZotfile(dir); err != nil {
154 t.Fatalf("mode %q rejected: %v", mode, err)
155 }
156 }
157 dir := writeTestZotfile(t, `{"zotfile":1,"name":"test","permissions":{"bash":{"mode":"all"}}}`)
158 if _, _, err := loadZotfile(dir); err == nil || !strings.Contains(err.Error(), "unsupported bash permission mode") {
159 t.Fatalf("unexpected error for all mode: %v", err)
160 }
161 dir = writeTestZotfile(t, `{"zotfile":1,"name":"test","permissions":{"bash":{"mode":"everything"}}}`)
162 if _, _, err := loadZotfile(dir); err == nil || !strings.Contains(err.Error(), "unsupported bash permission mode") {
163 t.Fatalf("unexpected error for unknown mode: %v", err)
164 }
165}
166
167func TestLoadZotfileRejectsUnsafeOrCollidingNames(t *testing.T) {
168 for _, name := range []string{"...", "Name", "two words", "a/b"} {
169 dir := writeTestZotfile(t, `{"zotfile":1,"name":"`+name+`"}`)
170 if _, _, err := loadZotfile(dir); err == nil {
171 t.Fatalf("unsafe manifest name %q was accepted", name)
172 }
173 }
174}
175
176func TestResolveZotfileRefUsesLocalFirstThenGitHubShorthand(t *testing.T) {
177 dir := t.TempDir()
178 t.Chdir(dir)
179 if err := os.Mkdir("local-agent", 0o700); err != nil {
180 t.Fatal(err)
181 }
182 if err := os.WriteFile("packed-agent.zot", []byte("test"), 0o600); err != nil {
183 t.Fatal(err)
184 }
185
186 tests := []struct {
187 ref string
188 want string
189 }{
190 {"local-agent", "local-agent"},
191 {"packed-agent", "packed-agent.zot"},
192 {"remote-agent", "remote-agent"},
193 {"missing.zot", "missing.zot"},
194 {"agents/remote-agent", "https://github.com/agents/remote-agent"},
195 {`agents\remote-agent`, "https://github.com/agents/remote-agent"},
196 {"frkr/zot-archify", "https://github.com/frkr/zot-archify"},
197 {"acme/agents/reviewer", "https://github.com/acme/agents/reviewer"},
198 {`acme\agents\reviewers\go`, "https://github.com/acme/agents/reviewers/go"},
199 {"./missing-agent", "./missing-agent"},
200 {"Remote-Agent", "Remote-Agent"},
201 {"https://github.com/acme/agents/example", "https://github.com/acme/agents/example"},
202 }
203 for _, tt := range tests {
204 got, err := resolveZotfileRef(tt.ref)
205 if err != nil {
206 t.Fatalf("resolve %q: %v", tt.ref, err)
207 }
208 if got != tt.want {
209 t.Errorf("resolve %q = %q, want %q", tt.ref, got, tt.want)
210 }
211 }
212}
213
214func TestLoadZotfileRejectsBundledExecutableExtension(t *testing.T) {
215 dir := writeTestZotfile(t, `{"zotfile":1,"name":"test"}`)
216 ext := filepath.Join(dir, "extensions", "bad")
217 if err := os.MkdirAll(ext, 0o700); err != nil {
218 t.Fatal(err)
219 }
220 if err := os.WriteFile(filepath.Join(ext, "extension.json"), []byte(`{"name":"bad"}`), 0o600); err != nil {
221 t.Fatal(err)
222 }
223 if _, _, err := loadZotfile(dir); err == nil || !strings.Contains(err.Error(), "cannot yet be confined") {
224 t.Fatalf("unexpected error: %v", err)
225 }
226}
227
228func TestCheckZotfileMinVersion(t *testing.T) {
229 var zf zotfileLoaded
230 zf.Manifest.Runtime.MinZot = "0.3.0"
231 if err := checkZotfileRequirements(zf, "0.2.75"); err == nil {
232 t.Fatal("old zot version accepted")
233 }
234 if err := checkZotfileRequirements(zf, "0.3.0"); err != nil {
235 t.Fatalf("minimum version rejected: %v", err)
236 }
237 if err := checkZotfileRequirements(zf, "not-a-version"); err == nil {
238 t.Fatal("invalid zot version accepted")
239 }
240
241 zf.Manifest.Runtime.MinZot = "not-a-version"
242 if err := checkZotfileRequirements(zf, "0.3.0"); err == nil {
243 t.Fatal("invalid min_zot requirement accepted")
244 }
245}
246
247func TestApplyZotfileModelRequirementsRejectsUnsupportedFields(t *testing.T) {
248 var m ZotfileManifest
249 m.Model.MinTier = "frontier"
250 if err := applyZotfileModelRequirements(&Args{}, m); err == nil {
251 t.Fatal("unsupported min_tier was ignored")
252 }
253 m.Model.MinTier = ""
254 m.Model.Requires = []string{"audio"}
255 if err := applyZotfileModelRequirements(&Args{}, m); err == nil {
256 t.Fatal("unsupported capability was ignored")
257 }
258}
259
260func TestUntarRejectsTraversalAndOversizedEntry(t *testing.T) {
261 makeTar := func(name string, size int64) []byte {
262 var buf bytes.Buffer
263 tw := tar.NewWriter(&buf)
264 if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: size}); err != nil {
265 t.Fatal(err)
266 }
267 _ = tw.Close()
268 return buf.Bytes()
269 }
270 if err := untar(bytes.NewReader(makeTar("../escape", 0)), t.TempDir()); err == nil {
271 t.Fatal("path traversal accepted")
272 }
273 if err := untar(bytes.NewReader(makeTar("large", maxZotfileEntrySize+1)), t.TempDir()); err == nil {
274 t.Fatal("oversized entry accepted")
275 }
276}
277
278func TestParseGitHubAgentURL(t *testing.T) {
279 tests := []struct {
280 input string
281 owner, repo, ref, dir string
282 }{
283 {"https://github.com/acme/agents/reviewer", "acme", "agents", "HEAD", "reviewer"},
284 {"https://github.com/acme/agents/tree/main/reviewer", "acme", "agents", "main", "reviewer"},
285 {"https://github.com/acme/agents/reviewers/go", "acme", "agents", "HEAD", "reviewers/go"},
286 {"https://github.com/acme/agent.git", "acme", "agent", "HEAD", ""},
287 }
288 for _, tt := range tests {
289 u, err := url.Parse(tt.input)
290 if err != nil {
291 t.Fatal(err)
292 }
293 owner, repo, ref, dir, err := parseGitHubAgentURL(u)
294 if err != nil {
295 t.Fatalf("parse %s: %v", tt.input, err)
296 }
297 if owner != tt.owner || repo != tt.repo || ref != tt.ref || dir != tt.dir {
298 t.Fatalf("parse %s = %q, %q, %q, %q", tt.input, owner, repo, ref, dir)
299 }
300 }
301}
302
303func TestLoadRemoteZotfileDownloadsTemporaryGitHubArchive(t *testing.T) {
304 var archive bytes.Buffer
305 gz := gzip.NewWriter(&archive)
306 tw := tar.NewWriter(gz)
307 files := map[string]string{
308 "agents-main/zot-maintenance/manifest.json": `{"zotfile":1,"name":"zot-maintenance"}`,
309 "agents-main/zot-maintenance/AGENT.md": "Maintain zot.",
310 }
311 for name, content := range files {
312 if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: int64(len(content))}); err != nil {
313 t.Fatal(err)
314 }
315 if _, err := tw.Write([]byte(content)); err != nil {
316 t.Fatal(err)
317 }
318 }
319 if err := tw.Close(); err != nil {
320 t.Fatal(err)
321 }
322 if err := gz.Close(); err != nil {
323 t.Fatal(err)
324 }
325
326 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
327 w.Header().Set("Content-Type", "application/gzip")
328 _, _ = w.Write(archive.Bytes())
329 }))
330 defer server.Close()
331 oldArchiveURL := githubArchiveURL
332 githubArchiveURL = func(_, _, _ string) string { return server.URL }
333 t.Cleanup(func() { githubArchiveURL = oldArchiveURL })
334
335 u, _ := url.Parse("https://github.com/acme/agents/zot-maintenance")
336 zf, cleanup, err := loadRemoteZotfile(u)
337 if err != nil {
338 t.Fatal(err)
339 }
340 if zf.Manifest.Name != "zot-maintenance" || !zf.Temp {
341 t.Fatalf("unexpected zotfile: %+v", zf)
342 }
343 root := filepath.Dir(filepath.Dir(zf.Dir))
344 cleanup()
345 if _, err := os.Stat(root); !os.IsNotExist(err) {
346 t.Fatalf("temporary checkout was not removed: %v", err)
347 }
348}
349
350func TestFormatZotfileConsentUsesThemeColors(t *testing.T) {
351 th := tui.Theme{FG: 1, Muted: 2, Accent: 3, Assistant: 4, Warning: 5}
352 zf := zotfileLoaded{}
353 zf.Manifest.Name = "reviewer"
354 zf.Manifest.Version = "1.2.3"
355 perms := tools.PermissionSet{}
356 perms.FS.Read = []string{"/workspace"}
357 perms.Bash.Mode = "ask"
358
359 got := formatZotfileConsent(zf, perms, th, true)
360 for _, want := range []string{
361 th.FG256(th.Assistant, "Agent"),
362 th.FG256(th.Accent, "reviewer@1.2.3"),
363 th.FG256(th.Muted, " fs read: "),
364 th.FG256(th.FG, "/workspace"),
365 th.FG256(th.Warning, "ask"),
366 th.FG256(th.Assistant, "Allow?"),
367 th.FG256(th.Muted, "[y/n]"),
368 } {
369 if !strings.Contains(got, want) {
370 t.Fatalf("themed consent missing %q:\n%s", want, got)
371 }
372 }
373}
374
375func TestPermissionSummaryShowsDeniedScopes(t *testing.T) {
376 got := permissionSummary(tools.PermissionSet{})
377 if !strings.Contains(got, "fs read: none") || !strings.Contains(got, "fs write: none") {
378 t.Fatalf("summary did not show denied scopes:\n%s", got)
379 }
380}
381
382func TestPeelLeadingYes(t *testing.T) {
383 yes, rest := peelLeadingYes([]string{"-y", "--yes", "agent", "-y"})
384 if !yes {
385 t.Fatal("expected yes")
386 }
387 if len(rest) != 2 || rest[0] != "agent" || rest[1] != "-y" {
388 t.Fatalf("rest = %v", rest)
389 }
390 yes, rest = peelLeadingYes([]string{"agent", "-y"})
391 if yes || len(rest) != 2 {
392 t.Fatalf("yes=%v rest=%v", yes, rest)
393 }
394}
395
396func TestConsentZotfileAutoYesWritesReceipt(t *testing.T) {
397 home := t.TempDir()
398 t.Setenv("ZOT_HOME", home)
399 zf := zotfileLoaded{Digest: "abc123", Manifest: ZotfileManifest{Name: "demo", Version: "1.0.0"}}
400 var perms tools.PermissionSet
401 perms.Bash.Mode = "none"
402 allowed, err := consentZotfile(zf, perms, true)
403 if err != nil || !allowed {
404 t.Fatalf("auto yes: allowed=%v err=%v", allowed, err)
405 }
406 path := filepath.Join(home, "agents", "demo", "consents", "abc123.json")
407 if _, err := os.Stat(path); err != nil {
408 t.Fatalf("consent receipt missing: %v", err)
409 }
410 // Second call without -y should use the receipt.
411 allowed, err = consentZotfile(zf, perms, false)
412 if err != nil || !allowed {
413 t.Fatalf("cached consent: allowed=%v err=%v", allowed, err)
414 }
415}
416
417func TestConsentZotfileAutoYesAskDoesNotCache(t *testing.T) {
418 home := t.TempDir()
419 t.Setenv("ZOT_HOME", home)
420 zf := zotfileLoaded{Digest: "askdig", Manifest: ZotfileManifest{Name: "ask-agent", Version: "1.0.0"}}
421 var perms tools.PermissionSet
422 perms.Bash.Mode = "ask"
423 if _, err := consentZotfile(zf, perms, true); err != nil {
424 t.Fatal(err)
425 }
426 path := filepath.Join(home, "agents", "ask-agent", "consents", "askdig.json")
427 if _, err := os.Stat(path); !os.IsNotExist(err) {
428 t.Fatalf("ask mode should not cache consent, err=%v", err)
429 }
430}