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 TestToolConfirmationRequestedEvent(t *testing.T) {
224 h := newHarness("confirmation-events")
225 received := make(chan Event, 1)
226 h.ext.On("tool_confirmation_requested", func(event Event) {
227 received <- event
228 })
229
230 go h.ext.Run()
231 h.handshake(t)
232 h.sendToExt(t, extproto.EventFromHost{
233 Type: "event",
234 Event: "tool_confirmation_requested",
235 ToolID: "call-1",
236 ToolName: "bash",
237 ToolPreview: "go test ./...",
238 })
239
240 select {
241 case event := <-received:
242 if event.Name != "tool_confirmation_requested" || event.ToolID != "call-1" || event.ToolName != "bash" || event.ToolPreview != "go test ./..." {
243 t.Fatalf("event = %+v", event)
244 }
245 case <-time.After(2 * time.Second):
246 t.Fatal("timeout waiting for confirmation event")
247 }
248 h.hostW.Close()
249}
250
251func TestDeferredToolRegistrationAndActivation(t *testing.T) {
252 h := newHarness("deferred-ext")
253 h.ext.DeferredTool("weather", "weather lookup", json.RawMessage(`{"type":"object"}`), func(json.RawMessage) ToolResult {
254 return TextResult("sunny")
255 })
256 h.ext.Tool("search_tools", "find tools", json.RawMessage(`{"type":"object"}`), func(json.RawMessage) ToolResult {
257 result := TextResult("enabled")
258 result.ActivateTools = []string{"weather"}
259 return result
260 })
261
262 go h.ext.Run()
263 if frame := h.next(t); frame.hdr.Type != "hello" {
264 t.Fatalf("expected hello, got %q", frame.hdr.Type)
265 }
266 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", ProtocolVersion: extproto.ProtocolVersion})
267
268 deferred := false
269 for {
270 frame := h.next(t)
271 if frame.hdr.Type == "ready" {
272 break
273 }
274 if frame.hdr.Type == "register_tool" {
275 var registration extproto.RegisterToolFromExt
276 if err := json.Unmarshal(frame.raw, ®istration); err != nil {
277 t.Fatal(err)
278 }
279 if registration.Name == "weather" {
280 deferred = registration.Deferred
281 }
282 }
283 }
284 if !deferred {
285 t.Fatal("weather registration was not deferred")
286 }
287
288 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "call-1", Name: "search_tools", Args: json.RawMessage(`{}`)})
289 frame := h.drainUntil(t, "tool_result")
290 var result extproto.ToolResultFromExt
291 if err := json.Unmarshal(frame.raw, &result); err != nil {
292 t.Fatal(err)
293 }
294 if len(result.ActivateTools) != 1 || result.ActivateTools[0] != "weather" {
295 t.Fatalf("activate_tools = %v", result.ActivateTools)
296 }
297 h.hostW.Close()
298}
299
300// TestBlockingToolWaitsForPanelKey is the core integration test for
301// the human-in-the-loop pattern: the tool handler opens a panel,
302// blocks on a channel, and only returns a tool_result after a key
303// event arrives.
304func TestBlockingToolWaitsForPanelKey(t *testing.T) {
305 h := newHarness("gate-ext")
306
307 const pid = "gate-panel"
308 const toolCallID = "tc-001"
309
310 approved := make(chan bool, 1)
311
312 h.ext.OnPanelKey(pid, func(key, text string) {
313 switch {
314 case key == "rune" && text == "y":
315 h.ext.ClosePanel(pid)
316 approved <- true
317 case key == "rune" && text == "n", key == "esc":
318 h.ext.ClosePanel(pid)
319 approved <- false
320 }
321 }, func() { approved <- false })
322
323 h.ext.Tool("gate", "needs approval",
324 json.RawMessage(`{"type":"object","properties":{}}`),
325 func(args json.RawMessage) ToolResult {
326 h.ext.OpenPanel(pid, "Approve?",
327 []string{" y approve", " n deny"}, "y/n")
328 if <-approved {
329 return TextResult("approved")
330 }
331 return TextErrorResult("denied")
332 })
333
334 go h.ext.Run()
335 h.handshake(t)
336
337 h.sendToExt(t, extproto.ToolCallFromHost{
338 Type: "tool_call", ID: toolCallID, Name: "gate",
339 Args: json.RawMessage(`{}`),
340 })
341
342 // Tool goroutine must open the panel before it can reply.
343 h.drainUntil(t, "open_panel")
344
345 // Send approval — tool should now unblock and emit tool_result.
346 h.sendToExt(t, extproto.PanelKeyFromHost{
347 Type: "panel_key", PanelID: pid, Key: "rune", Text: "y",
348 })
349
350 f := h.drainUntil(t, "tool_result")
351 var tr extproto.ToolResultFromExt
352 if err := json.Unmarshal(f.raw, &tr); err != nil {
353 t.Fatalf("unmarshal tool_result: %v", err)
354 }
355 if tr.ID != toolCallID {
356 t.Errorf("tool_result id: want %q, got %q", toolCallID, tr.ID)
357 }
358 if tr.IsError {
359 t.Errorf("expected success, got is_error=true")
360 }
361 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "approved") {
362 t.Errorf("expected 'approved' in content, got %+v", tr.Content)
363 }
364
365 h.hostW.Close()
366}
367
368// TestBlockingToolDenied mirrors TestBlockingToolWaitsForPanelKey but
369// sends "n" so the tool returns an error result.
370func TestBlockingToolDenied(t *testing.T) {
371 h := newHarness("gate-ext-deny")
372
373 const pid = "deny-panel"
374 const toolCallID = "tc-002"
375
376 approved := make(chan bool, 1)
377
378 h.ext.OnPanelKey(pid, func(key, text string) {
379 switch {
380 case key == "rune" && text == "y":
381 h.ext.ClosePanel(pid)
382 approved <- true
383 case key == "rune" && text == "n", key == "esc":
384 h.ext.ClosePanel(pid)
385 approved <- false
386 }
387 }, func() { approved <- false })
388
389 h.ext.Tool("gate2", "needs approval",
390 json.RawMessage(`{"type":"object","properties":{}}`),
391 func(args json.RawMessage) ToolResult {
392 h.ext.OpenPanel(pid, "Approve?", []string{"y/n"}, "")
393 if <-approved {
394 return TextResult("approved")
395 }
396 return TextErrorResult("denied")
397 })
398
399 go h.ext.Run()
400 h.handshake(t)
401
402 h.sendToExt(t, extproto.ToolCallFromHost{
403 Type: "tool_call", ID: toolCallID, Name: "gate2",
404 Args: json.RawMessage(`{}`),
405 })
406
407 h.drainUntil(t, "open_panel")
408
409 h.sendToExt(t, extproto.PanelKeyFromHost{
410 Type: "panel_key", PanelID: pid, Key: "rune", Text: "n",
411 })
412
413 f := h.drainUntil(t, "tool_result")
414 var tr extproto.ToolResultFromExt
415 if err := json.Unmarshal(f.raw, &tr); err != nil {
416 t.Fatalf("unmarshal tool_result: %v", err)
417 }
418 if !tr.IsError {
419 t.Errorf("expected is_error=true on denial")
420 }
421 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "denied") {
422 t.Errorf("expected 'denied' in content, got %+v", tr.Content)
423 }
424
425 h.hostW.Close()
426}