1package ext
2
3import (
4 "bufio"
5 "encoding/json"
6 "io"
7 "strings"
8 "testing"
9 "time"
10
11 "github.com/patriceckhart/zot/packages/agent/extproto"
12)
13
14// ---------- test harness ----------
15
16// extHarness wires an Extension to io.Pipe pairs so a test can play
17// the role of the host: write host→ext frames, read ext→host frames.
18// The scanner runs in a permanent background goroutine and delivers
19// frames over a buffered channel, avoiding the deadlock that would
20// occur if the test goroutine alternated between writing and reading
21// a synchronous pipe.
22type extHarness struct {
23 ext *Extension
24 hostW *io.PipeWriter // test writes here → ext reads as stdin
25 frames chan rawFrame // ext→host frames delivered here
26}
27
28type rawFrame struct {
29 hdr extproto.Frame
30 raw []byte
31}
32
33func newHarness(name string) *extHarness {
34 extStdinR, extStdinW := io.Pipe()
35 extStdoutR, extStdoutW := io.Pipe()
36
37 e := New(name, "0.0.0-test")
38 e.in = extStdinR
39 e.out = extStdoutW
40 e.stderr = io.Discard
41
42 h := &extHarness{
43 ext: e,
44 hostW: extStdinW,
45 frames: make(chan rawFrame, 64),
46 }
47
48 // Background reader: scan ext's stdout and push every frame into
49 // the channel so the test goroutine never needs to block on the pipe.
50 go func() {
51 scanner := bufio.NewScanner(extStdoutR)
52 scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
53 for scanner.Scan() {
54 b := scanner.Bytes()
55 cp := make([]byte, len(b))
56 copy(cp, b)
57 var f extproto.Frame
58 json.Unmarshal(cp, &f)
59 h.frames <- rawFrame{f, cp}
60 }
61 close(h.frames)
62 }()
63
64 return h
65}
66
67// next returns the next frame, timing out after 2 s.
68func (h *extHarness) next(t *testing.T) rawFrame {
69 t.Helper()
70 select {
71 case f, ok := <-h.frames:
72 if !ok {
73 t.Fatal("frame channel closed (ext stdout EOF)")
74 }
75 return f
76 case <-time.After(2 * time.Second):
77 t.Fatal("timeout waiting for frame from extension")
78 return rawFrame{}
79 }
80}
81
82// drainUntil reads frames until one with type == want arrives.
83func (h *extHarness) drainUntil(t *testing.T, want string) rawFrame {
84 t.Helper()
85 deadline := time.NewTimer(2 * time.Second)
86 defer deadline.Stop()
87 for {
88 select {
89 case f, ok := <-h.frames:
90 if !ok {
91 t.Fatalf("frame channel closed before seeing %q", want)
92 }
93 if f.hdr.Type == want {
94 return f
95 }
96 case <-deadline.C:
97 t.Fatalf("timeout waiting for frame type %q", want)
98 return rawFrame{}
99 }
100 }
101}
102
103// sendToExt writes a host→ext frame.
104func (h *extHarness) sendToExt(t *testing.T, v any) {
105 t.Helper()
106 b, err := extproto.Encode(v)
107 if err != nil {
108 t.Fatalf("encode: %v", err)
109 }
110 if _, err := h.hostW.Write(b); err != nil {
111 t.Fatalf("write to ext: %v", err)
112 }
113}
114
115// handshake performs the hello / hello_ack exchange and drains frames
116// until "ready".
117func (h *extHarness) handshake(t *testing.T) {
118 t.Helper()
119 f := h.next(t)
120 if f.hdr.Type != "hello" {
121 t.Fatalf("expected hello, got %q", f.hdr.Type)
122 }
123 h.sendToExt(t, extproto.HelloAckFromHost{
124 Type: "hello_ack",
125 ProtocolVersion: extproto.ProtocolVersion,
126 ZotVersion: "0.0.0-test",
127 Provider: "anthropic",
128 Model: "claude-test",
129 })
130 for {
131 f := h.next(t)
132 if f.hdr.Type == "ready" {
133 return
134 }
135 }
136}
137
138// ---------- tests ----------
139
140// TestOnHelloCanRegisterUsingHostInfo checks that extensions can use
141// host metadata before announcing their initial registrations.
142func TestOnHelloCanRegisterUsingHostInfo(t *testing.T) {
143 h := newHarness("cwd-ext")
144
145 h.ext.OnHello(func(host HostInfo) {
146 h.ext.Command("cwd", "show cwd", func(args string) Response {
147 return Display(host.CWD)
148 })
149 })
150
151 go h.ext.Run()
152
153 f := h.next(t)
154 if f.hdr.Type != "hello" {
155 t.Fatalf("expected hello, got %q", f.hdr.Type)
156 }
157 h.sendToExt(t, extproto.HelloAckFromHost{
158 Type: "hello_ack",
159 ProtocolVersion: extproto.ProtocolVersion,
160 ZotVersion: "0.0.0-test",
161 Provider: "anthropic",
162 Model: "claude-test",
163 CWD: "/tmp/project",
164 })
165
166 var sawCommand bool
167 for {
168 f := h.next(t)
169 if f.hdr.Type == "register_command" {
170 var rc extproto.RegisterCommandFromExt
171 if err := json.Unmarshal(f.raw, &rc); err != nil {
172 t.Fatalf("unmarshal register_command: %v", err)
173 }
174 if rc.Name == "cwd" {
175 sawCommand = true
176 }
177 }
178 if f.hdr.Type == "ready" {
179 break
180 }
181 }
182 if !sawCommand {
183 t.Fatal("OnHello-registered command was not announced before ready")
184 }
185 if got := h.ext.Host().CWD; got != "/tmp/project" {
186 t.Fatalf("Host().CWD = %q, want /tmp/project", got)
187 }
188
189 h.hostW.Close()
190}
191
192// TestOpenPanelEmitsCorrectFrame checks that e.OpenPanel sends a
193// well-formed open_panel frame with the correct PanelSpec fields.
194func TestOpenPanelEmitsCorrectFrame(t *testing.T) {
195 h := newHarness("test-ext")
196 go h.ext.Run()
197 h.handshake(t)
198
199 go h.ext.OpenPanel("my-panel", "My Title", []string{"line a", "line b"}, "esc close")
200
201 f := h.drainUntil(t, "open_panel")
202
203 var op extproto.OpenPanelFromExt
204 if err := json.Unmarshal(f.raw, &op); err != nil {
205 t.Fatalf("unmarshal open_panel: %v", err)
206 }
207 if op.Panel.ID != "my-panel" {
208 t.Errorf("panel id: want %q, got %q", "my-panel", op.Panel.ID)
209 }
210 if op.Panel.Title != "My Title" {
211 t.Errorf("panel title: want %q, got %q", "My Title", op.Panel.Title)
212 }
213 if len(op.Panel.Lines) != 2 || op.Panel.Lines[0] != "line a" || op.Panel.Lines[1] != "line b" {
214 t.Errorf("panel lines: got %v", op.Panel.Lines)
215 }
216 if op.Panel.Footer != "esc close" {
217 t.Errorf("panel footer: want %q, got %q", "esc close", op.Panel.Footer)
218 }
219
220 h.hostW.Close()
221}
222
223func TestDeferredToolRegistrationAndActivation(t *testing.T) {
224 h := newHarness("deferred-ext")
225 h.ext.DeferredTool("weather", "weather lookup", json.RawMessage(`{"type":"object"}`), func(json.RawMessage) ToolResult {
226 return TextResult("sunny")
227 })
228 h.ext.Tool("search_tools", "find tools", json.RawMessage(`{"type":"object"}`), func(json.RawMessage) ToolResult {
229 result := TextResult("enabled")
230 result.ActivateTools = []string{"weather"}
231 return result
232 })
233
234 go h.ext.Run()
235 if frame := h.next(t); frame.hdr.Type != "hello" {
236 t.Fatalf("expected hello, got %q", frame.hdr.Type)
237 }
238 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", ProtocolVersion: extproto.ProtocolVersion})
239
240 deferred := false
241 for {
242 frame := h.next(t)
243 if frame.hdr.Type == "ready" {
244 break
245 }
246 if frame.hdr.Type == "register_tool" {
247 var registration extproto.RegisterToolFromExt
248 if err := json.Unmarshal(frame.raw, ®istration); err != nil {
249 t.Fatal(err)
250 }
251 if registration.Name == "weather" {
252 deferred = registration.Deferred
253 }
254 }
255 }
256 if !deferred {
257 t.Fatal("weather registration was not deferred")
258 }
259
260 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "call-1", Name: "search_tools", Args: json.RawMessage(`{}`)})
261 frame := h.drainUntil(t, "tool_result")
262 var result extproto.ToolResultFromExt
263 if err := json.Unmarshal(frame.raw, &result); err != nil {
264 t.Fatal(err)
265 }
266 if len(result.ActivateTools) != 1 || result.ActivateTools[0] != "weather" {
267 t.Fatalf("activate_tools = %v", result.ActivateTools)
268 }
269 h.hostW.Close()
270}
271
272// TestBlockingToolWaitsForPanelKey is the core integration test for
273// the human-in-the-loop pattern: the tool handler opens a panel,
274// blocks on a channel, and only returns a tool_result after a key
275// event arrives.
276func TestBlockingToolWaitsForPanelKey(t *testing.T) {
277 h := newHarness("gate-ext")
278
279 const pid = "gate-panel"
280 const toolCallID = "tc-001"
281
282 approved := make(chan bool, 1)
283
284 h.ext.OnPanelKey(pid, func(key, text string) {
285 switch {
286 case key == "rune" && text == "y":
287 h.ext.ClosePanel(pid)
288 approved <- true
289 case key == "rune" && text == "n", key == "esc":
290 h.ext.ClosePanel(pid)
291 approved <- false
292 }
293 }, func() { approved <- false })
294
295 h.ext.Tool("gate", "needs approval",
296 json.RawMessage(`{"type":"object","properties":{}}`),
297 func(args json.RawMessage) ToolResult {
298 h.ext.OpenPanel(pid, "Approve?",
299 []string{" y approve", " n deny"}, "y/n")
300 if <-approved {
301 return TextResult("approved")
302 }
303 return TextErrorResult("denied")
304 })
305
306 go h.ext.Run()
307 h.handshake(t)
308
309 h.sendToExt(t, extproto.ToolCallFromHost{
310 Type: "tool_call", ID: toolCallID, Name: "gate",
311 Args: json.RawMessage(`{}`),
312 })
313
314 // Tool goroutine must open the panel before it can reply.
315 h.drainUntil(t, "open_panel")
316
317 // Send approval — tool should now unblock and emit tool_result.
318 h.sendToExt(t, extproto.PanelKeyFromHost{
319 Type: "panel_key", PanelID: pid, Key: "rune", Text: "y",
320 })
321
322 f := h.drainUntil(t, "tool_result")
323 var tr extproto.ToolResultFromExt
324 if err := json.Unmarshal(f.raw, &tr); err != nil {
325 t.Fatalf("unmarshal tool_result: %v", err)
326 }
327 if tr.ID != toolCallID {
328 t.Errorf("tool_result id: want %q, got %q", toolCallID, tr.ID)
329 }
330 if tr.IsError {
331 t.Errorf("expected success, got is_error=true")
332 }
333 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "approved") {
334 t.Errorf("expected 'approved' in content, got %+v", tr.Content)
335 }
336
337 h.hostW.Close()
338}
339
340// TestBlockingToolDenied mirrors TestBlockingToolWaitsForPanelKey but
341// sends "n" so the tool returns an error result.
342func TestBlockingToolDenied(t *testing.T) {
343 h := newHarness("gate-ext-deny")
344
345 const pid = "deny-panel"
346 const toolCallID = "tc-002"
347
348 approved := make(chan bool, 1)
349
350 h.ext.OnPanelKey(pid, func(key, text string) {
351 switch {
352 case key == "rune" && text == "y":
353 h.ext.ClosePanel(pid)
354 approved <- true
355 case key == "rune" && text == "n", key == "esc":
356 h.ext.ClosePanel(pid)
357 approved <- false
358 }
359 }, func() { approved <- false })
360
361 h.ext.Tool("gate2", "needs approval",
362 json.RawMessage(`{"type":"object","properties":{}}`),
363 func(args json.RawMessage) ToolResult {
364 h.ext.OpenPanel(pid, "Approve?", []string{"y/n"}, "")
365 if <-approved {
366 return TextResult("approved")
367 }
368 return TextErrorResult("denied")
369 })
370
371 go h.ext.Run()
372 h.handshake(t)
373
374 h.sendToExt(t, extproto.ToolCallFromHost{
375 Type: "tool_call", ID: toolCallID, Name: "gate2",
376 Args: json.RawMessage(`{}`),
377 })
378
379 h.drainUntil(t, "open_panel")
380
381 h.sendToExt(t, extproto.PanelKeyFromHost{
382 Type: "panel_key", PanelID: pid, Key: "rune", Text: "n",
383 })
384
385 f := h.drainUntil(t, "tool_result")
386 var tr extproto.ToolResultFromExt
387 if err := json.Unmarshal(f.raw, &tr); err != nil {
388 t.Fatalf("unmarshal tool_result: %v", err)
389 }
390 if !tr.IsError {
391 t.Errorf("expected is_error=true on denial")
392 }
393 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "denied") {
394 t.Errorf("expected 'denied' in content, got %+v", tr.Content)
395 }
396
397 h.hostW.Close()
398}