1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "os"
7
8 "github.com/patriceckhart/zot/packages/agent/extproto"
9 "path/filepath"
10 "runtime"
11 "strings"
12 "sync"
13 "testing"
14 "time"
15)
16
17// stubHooks records every callback so the test can assert on them.
18type stubHooks struct {
19 mu sync.Mutex
20 notifies []string
21 displays []string
22 submits []string
23 submitSlash []string
24 clearNotes []string
25 panels []extproto.PanelSpec
26 panelExts []string
27}
28
29func (s *stubHooks) Notify(name, level, message string) {
30 s.mu.Lock()
31 defer s.mu.Unlock()
32 s.notifies = append(s.notifies, name+":"+level+":"+message)
33}
34func (s *stubHooks) Submit(text string) {
35 s.mu.Lock()
36 defer s.mu.Unlock()
37 s.submits = append(s.submits, text)
38}
39func (s *stubHooks) SubmitSlash(text string) {
40 s.mu.Lock()
41 defer s.mu.Unlock()
42 s.submitSlash = append(s.submitSlash, text)
43}
44func (s *stubHooks) Insert(string) {}
45func (s *stubHooks) Display(name, text string) {
46 s.mu.Lock()
47 defer s.mu.Unlock()
48 s.displays = append(s.displays, name+":"+text)
49}
50func (s *stubHooks) ClearNotes(name string) {
51 s.mu.Lock()
52 defer s.mu.Unlock()
53 s.clearNotes = append(s.clearNotes, name)
54}
55func (s *stubHooks) OpenPanel(extName string, spec extproto.PanelSpec) {
56 s.mu.Lock()
57 defer s.mu.Unlock()
58 s.panelExts = append(s.panelExts, extName)
59 s.panels = append(s.panels, spec)
60}
61func (s *stubHooks) UpdatePanel(string, string, string, []string, string) {}
62func (s *stubHooks) ClosePanel(string, string) {}
63
64// writeMockExtension creates a minimal extension on disk that uses a
65// shell script (or batch file on windows) to drive the protocol. The
66// script reads commands from stdin and emits hard-coded responses,
67// exercising the manager's spawn/handshake/dispatch path without
68// needing the SDK.
69func writeMockExtension(t *testing.T, root string) {
70 t.Helper()
71 if runtime.GOOS == "windows" {
72 t.Skip("mock extension uses /bin/sh; skip on windows")
73 }
74
75 dir := filepath.Join(root, "mock")
76 if err := os.MkdirAll(dir, 0o755); err != nil {
77 t.Fatal(err)
78 }
79
80 // Shell script: emit hello, read frames, respond. Reads until
81 // stdin closes; tail's -F keeps the pipe alive long enough for
82 // the manager to send command_invoked.
83 script := `#!/bin/sh
84printf '%s\n' '{"type":"hello","name":"mock","version":"0.0.1","capabilities":["commands"]}'
85printf '%s\n' '{"type":"register_command","name":"ping","description":"ping/pong"}'
86while IFS= read -r line; do
87 case "$line" in
88 *'"type":"command_invoked"'*)
89 id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([^"]*\)".*/\1/p')
90 printf '%s\n' "{\"type\":\"command_response\",\"id\":\"$id\",\"action\":\"display\",\"display\":\"pong\"}"
91 ;;
92 *'"type":"shutdown"'*)
93 printf '%s\n' '{"type":"shutdown_ack"}'
94 exit 0
95 ;;
96 esac
97done
98`
99 scriptPath := filepath.Join(dir, "run.sh")
100 if err := os.WriteFile(scriptPath, []byte(script), 0o755); err != nil {
101 t.Fatal(err)
102 }
103
104 manifest := map[string]any{
105 "name": "mock",
106 "version": "0.0.1",
107 "exec": "./run.sh",
108 }
109 mfb, _ := json.Marshal(manifest)
110 if err := os.WriteFile(filepath.Join(dir, "extension.json"), mfb, 0o644); err != nil {
111 t.Fatal(err)
112 }
113}
114
115func TestDiscoverLoadsThemeOnlyExtension(t *testing.T) {
116 tmp := t.TempDir()
117 extDir := filepath.Join(tmp, "extensions", "theme-only")
118 if err := os.MkdirAll(extDir, 0o755); err != nil {
119 t.Fatal(err)
120 }
121 manifest := `{"name":"theme-only","version":"1.0.0","description":"theme only"}`
122 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), []byte(manifest), 0o644); err != nil {
123 t.Fatal(err)
124 }
125 theme := `{"name":"Theme Only","description":"theme from extension","colors":{"dark":{"accent":204}}}`
126 if err := os.WriteFile(filepath.Join(extDir, "theme.json"), []byte(theme), 0o644); err != nil {
127 t.Fatal(err)
128 }
129
130 mgr := New(tmp, "", "0.0.0-test", "anthropic", "claude-opus-4-7", nil)
131 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
132 t.Fatalf("discover errors: %v", errs)
133 }
134 defer mgr.Stop(10 * time.Millisecond)
135
136 opts := mgr.ThemeOptions()
137 if len(opts) != 1 {
138 t.Fatalf("theme options = %d, want 1", len(opts))
139 }
140 if opts[0].Label != "Theme Only" || opts[0].Path != filepath.Join(extDir, "theme.json") {
141 t.Fatalf("unexpected theme option: %#v", opts[0])
142 }
143 if !strings.Contains(opts[0].Description, "from extension theme-only") {
144 t.Fatalf("description missing extension source: %q", opts[0].Description)
145 }
146}
147
148func TestManagerSpawnAndInvoke(t *testing.T) {
149 tmp := t.TempDir()
150 extRoot := filepath.Join(tmp, "extensions")
151 if err := os.MkdirAll(extRoot, 0o755); err != nil {
152 t.Fatal(err)
153 }
154 writeMockExtension(t, extRoot)
155
156 hooks := &stubHooks{}
157 mgr := New(tmp, "", "0.0.0-test", "anthropic", "claude-opus-4-7", hooks)
158
159 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
160 t.Fatalf("discover errors: %v", errs)
161 }
162 defer mgr.Stop(2 * time.Second)
163
164 // Give the extension a beat to send register_command frames after
165 // the hello handshake.
166 time.Sleep(150 * time.Millisecond)
167
168 cmds := mgr.Commands()
169 found := false
170 for _, c := range cmds {
171 if c.Name == "ping" {
172 found = true
173 break
174 }
175 }
176 if !found {
177 t.Fatalf("expected command 'ping', got %#v", cmds)
178 }
179 if !mgr.HasCommand("ping") {
180 t.Fatal("HasCommand(\"ping\") = false")
181 }
182
183 resp, err := mgr.Invoke(context.Background(), "ping", "", 2*time.Second)
184 if err != nil {
185 t.Fatalf("invoke: %v", err)
186 }
187 if resp.Action != "display" {
188 t.Errorf("expected action=display, got %q", resp.Action)
189 }
190 if resp.Display != "pong" {
191 t.Errorf("expected display=pong, got %q", resp.Display)
192 }
193}
194
195func TestDiagnosticsReportMalformedFramesAndConflicts(t *testing.T) {
196 if runtime.GOOS == "windows" {
197 t.Skip("mock extension uses /bin/sh; skip on windows")
198 }
199
200 tmp := t.TempDir()
201 extRoot := filepath.Join(tmp, "extensions")
202 writeDiagExtension := func(name, script string) {
203 t.Helper()
204 dir := filepath.Join(extRoot, name)
205 if err := os.MkdirAll(dir, 0o755); err != nil {
206 t.Fatal(err)
207 }
208 if err := os.WriteFile(filepath.Join(dir, "run.sh"), []byte(script), 0o755); err != nil {
209 t.Fatal(err)
210 }
211 mfb, _ := json.Marshal(map[string]any{"name": name, "exec": "./run.sh"})
212 if err := os.WriteFile(filepath.Join(dir, "extension.json"), mfb, 0o644); err != nil {
213 t.Fatal(err)
214 }
215 }
216
217 writeDiagExtension("a-first", `#!/bin/sh
218printf '%s\n' '{"type":"hello","name":"a-first","version":"0.1","capabilities":["tools"]}'
219printf '%s\n' '{"type":"register_tool","name":"shared","description":"first","schema":{"type":"object"}}'
220printf '%s\n' '{"type":"register_tool","name":"broken","description":"bad","schema":'
221printf '%s\n' '{"type":"ready"}'
222while IFS= read -r line; do
223 case "$line" in
224 *'"type":"shutdown"'*) exit 0 ;;
225 esac
226done
227`)
228 writeDiagExtension("b-second", `#!/bin/sh
229printf '%s\n' '{"type":"hello","name":"b-second","version":"0.1","capabilities":["tools"]}'
230printf '%s\n' '{"type":"register_tool","name":"shared","description":"second","schema":{"type":"object"}}'
231printf '%s\n' '{"type":"ready"}'
232while IFS= read -r line; do
233 case "$line" in
234 *'"type":"shutdown"'*) exit 0 ;;
235 esac
236done
237`)
238
239 mgr := New(tmp, "", "0.0.0-test", "anthropic", "claude-opus-4-7", &stubHooks{})
240 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
241 t.Fatalf("discover errors: %v", errs)
242 }
243 defer mgr.Stop(2 * time.Second)
244 mgr.WaitForReady(2 * time.Second)
245
246 diags := mgr.Diagnostics()
247 byName := map[string]ExtensionDiagnostic{}
248 for _, d := range diags {
249 byName[d.Name] = d
250 }
251
252 first := byName["a-first"]
253 if len(first.Messages) == 0 || !strings.Contains(strings.Join(first.Messages, "\n"), "malformed json frame") {
254 t.Fatalf("expected malformed-frame diagnostic, got %#v", first.Messages)
255 }
256
257 var shadowedTool bool
258 var conflictMessage bool
259 for _, d := range diags {
260 if strings.Contains(strings.Join(d.Messages, "\n"), "conflicts with another extension") {
261 conflictMessage = true
262 }
263 for _, tool := range d.Tools {
264 if tool.Name == "shared" && !tool.Active {
265 shadowedTool = true
266 }
267 }
268 }
269 if !shadowedTool {
270 t.Fatalf("expected one shared tool registration to be inactive, got %#v", diags)
271 }
272 if !conflictMessage {
273 t.Fatalf("expected conflict diagnostic, got %#v", diags)
274 }
275}
276
277// TestSpontaneousSubmit verifies that an extension can submit a
278// non-empty prompt outside of any command response.
279func TestSpontaneousSubmit(t *testing.T) {
280 if runtime.GOOS == "windows" {
281 t.Skip("mock extension uses /bin/sh; skip on windows")
282 }
283
284 tmp := t.TempDir()
285 extDir := filepath.Join(tmp, "extensions", "submit-mock")
286 if err := os.MkdirAll(extDir, 0o755); err != nil {
287 t.Fatal(err)
288 }
289
290 script := `#!/bin/sh
291printf '%s\n' '{"type":"hello","name":"submit-mock","version":"0.1","capabilities":["submit"]}'
292printf '%s\n' '{"type":"ready"}'
293printf '%s\n' '{"type":"submit","text":" explain this repository briefly "}'
294printf '%s\n' '{"type":"submit","text":" "}'
295while IFS= read -r line; do
296 case "$line" in
297 *'"type":"shutdown"'*)
298 printf '%s\n' '{"type":"shutdown_ack"}'
299 exit 0
300 ;;
301 esac
302done
303`
304 if err := os.WriteFile(filepath.Join(extDir, "run.sh"), []byte(script), 0o755); err != nil {
305 t.Fatal(err)
306 }
307 mfb, _ := json.Marshal(map[string]any{"name": "submit-mock", "exec": "./run.sh"})
308 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), mfb, 0o644); err != nil {
309 t.Fatal(err)
310 }
311
312 hooks := &stubHooks{}
313 mgr := New(tmp, "", "0.0.0-test", "anthropic", "claude-opus-4-7", hooks)
314 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
315 t.Fatalf("discover errors: %v", errs)
316 }
317 defer mgr.Stop(2 * time.Second)
318
319 deadline := time.Now().Add(2 * time.Second)
320 for time.Now().Before(deadline) {
321 hooks.mu.Lock()
322 n := len(hooks.submits)
323 hooks.mu.Unlock()
324 if n > 0 {
325 break
326 }
327 time.Sleep(20 * time.Millisecond)
328 }
329
330 hooks.mu.Lock()
331 defer hooks.mu.Unlock()
332 if len(hooks.submits) != 1 {
333 t.Fatalf("Submit calls = %#v, want one non-empty call", hooks.submits)
334 }
335 if hooks.submits[0] != "explain this repository briefly" {
336 t.Fatalf("Submit text = %q", hooks.submits[0])
337 }
338}
339
340// TestSpontaneousOpenPanel verifies that an extension sending an
341// open_panel frame outside of any command response causes the manager
342// to call hooks.OpenPanel with the correct PanelSpec fields.
343func TestSpontaneousOpenPanel(t *testing.T) {
344 if runtime.GOOS == "windows" {
345 t.Skip("mock extension uses /bin/sh; skip on windows")
346 }
347
348 tmp := t.TempDir()
349 extDir := filepath.Join(tmp, "extensions", "panel-mock")
350 if err := os.MkdirAll(extDir, 0o755); err != nil {
351 t.Fatal(err)
352 }
353
354 // Extension emits hello + ready, then immediately fires a
355 // spontaneous open_panel, then waits for shutdown.
356 script := `#!/bin/sh
357printf '%s\n' '{"type":"hello","name":"panel-mock","version":"0.1","capabilities":["panels"]}'
358printf '%s\n' '{"type":"ready"}'
359printf '%s\n' '{"type":"open_panel","panel":{"id":"test-panel","title":"Hello Panel","lines":["line one","line two"],"footer":"esc close"}}'
360while IFS= read -r line; do
361 case "$line" in
362 *'"type":"shutdown"'*)
363 printf '%s\n' '{"type":"shutdown_ack"}'
364 exit 0
365 ;;
366 esac
367done
368`
369 if err := os.WriteFile(filepath.Join(extDir, "run.sh"), []byte(script), 0o755); err != nil {
370 t.Fatal(err)
371 }
372 mfb, _ := json.Marshal(map[string]any{"name": "panel-mock", "exec": "./run.sh"})
373 if err := os.WriteFile(filepath.Join(extDir, "extension.json"), mfb, 0o644); err != nil {
374 t.Fatal(err)
375 }
376
377 hooks := &stubHooks{}
378 mgr := New(tmp, "", "0.0.0-test", "anthropic", "claude-opus-4-7", hooks)
379 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
380 t.Fatalf("discover errors: %v", errs)
381 }
382 defer mgr.Stop(2 * time.Second)
383
384 // Give the extension time to flush its open_panel frame.
385 deadline := time.Now().Add(2 * time.Second)
386 for time.Now().Before(deadline) {
387 hooks.mu.Lock()
388 n := len(hooks.panels)
389 hooks.mu.Unlock()
390 if n > 0 {
391 break
392 }
393 time.Sleep(20 * time.Millisecond)
394 }
395
396 hooks.mu.Lock()
397 defer hooks.mu.Unlock()
398
399 if len(hooks.panels) == 0 {
400 t.Fatal("hooks.OpenPanel was never called")
401 }
402 spec := hooks.panels[0]
403 if spec.ID != "test-panel" {
404 t.Errorf("panel id: want %q, got %q", "test-panel", spec.ID)
405 }
406 if spec.Title != "Hello Panel" {
407 t.Errorf("panel title: want %q, got %q", "Hello Panel", spec.Title)
408 }
409 if len(spec.Lines) != 2 || spec.Lines[0] != "line one" || spec.Lines[1] != "line two" {
410 t.Errorf("panel lines: want [line one line two], got %v", spec.Lines)
411 }
412 if spec.Footer != "esc close" {
413 t.Errorf("panel footer: want %q, got %q", "esc close", spec.Footer)
414 }
415 if hooks.panelExts[0] != "panel-mock" {
416 t.Errorf("ext name: want %q, got %q", "panel-mock", hooks.panelExts[0])
417 }
418}