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
223// TestBlockingToolWaitsForPanelKey is the core integration test for
224// the human-in-the-loop pattern: the tool handler opens a panel,
225// blocks on a channel, and only returns a tool_result after a key
226// event arrives.
227func TestBlockingToolWaitsForPanelKey(t *testing.T) {
228 h := newHarness("gate-ext")
229
230 const pid = "gate-panel"
231 const toolCallID = "tc-001"
232
233 approved := make(chan bool, 1)
234
235 h.ext.OnPanelKey(pid, func(key, text string) {
236 switch {
237 case key == "rune" && text == "y":
238 h.ext.ClosePanel(pid)
239 approved <- true
240 case key == "rune" && text == "n", key == "esc":
241 h.ext.ClosePanel(pid)
242 approved <- false
243 }
244 }, func() { approved <- false })
245
246 h.ext.Tool("gate", "needs approval",
247 json.RawMessage(`{"type":"object","properties":{}}`),
248 func(args json.RawMessage) ToolResult {
249 h.ext.OpenPanel(pid, "Approve?",
250 []string{" y approve", " n deny"}, "y/n")
251 if <-approved {
252 return TextResult("approved")
253 }
254 return TextErrorResult("denied")
255 })
256
257 go h.ext.Run()
258 h.handshake(t)
259
260 h.sendToExt(t, extproto.ToolCallFromHost{
261 Type: "tool_call", ID: toolCallID, Name: "gate",
262 Args: json.RawMessage(`{}`),
263 })
264
265 // Tool goroutine must open the panel before it can reply.
266 h.drainUntil(t, "open_panel")
267
268 // Send approval — tool should now unblock and emit tool_result.
269 h.sendToExt(t, extproto.PanelKeyFromHost{
270 Type: "panel_key", PanelID: pid, Key: "rune", Text: "y",
271 })
272
273 f := h.drainUntil(t, "tool_result")
274 var tr extproto.ToolResultFromExt
275 if err := json.Unmarshal(f.raw, &tr); err != nil {
276 t.Fatalf("unmarshal tool_result: %v", err)
277 }
278 if tr.ID != toolCallID {
279 t.Errorf("tool_result id: want %q, got %q", toolCallID, tr.ID)
280 }
281 if tr.IsError {
282 t.Errorf("expected success, got is_error=true")
283 }
284 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "approved") {
285 t.Errorf("expected 'approved' in content, got %+v", tr.Content)
286 }
287
288 h.hostW.Close()
289}
290
291// TestBlockingToolDenied mirrors TestBlockingToolWaitsForPanelKey but
292// sends "n" so the tool returns an error result.
293func TestBlockingToolDenied(t *testing.T) {
294 h := newHarness("gate-ext-deny")
295
296 const pid = "deny-panel"
297 const toolCallID = "tc-002"
298
299 approved := make(chan bool, 1)
300
301 h.ext.OnPanelKey(pid, func(key, text string) {
302 switch {
303 case key == "rune" && text == "y":
304 h.ext.ClosePanel(pid)
305 approved <- true
306 case key == "rune" && text == "n", key == "esc":
307 h.ext.ClosePanel(pid)
308 approved <- false
309 }
310 }, func() { approved <- false })
311
312 h.ext.Tool("gate2", "needs approval",
313 json.RawMessage(`{"type":"object","properties":{}}`),
314 func(args json.RawMessage) ToolResult {
315 h.ext.OpenPanel(pid, "Approve?", []string{"y/n"}, "")
316 if <-approved {
317 return TextResult("approved")
318 }
319 return TextErrorResult("denied")
320 })
321
322 go h.ext.Run()
323 h.handshake(t)
324
325 h.sendToExt(t, extproto.ToolCallFromHost{
326 Type: "tool_call", ID: toolCallID, Name: "gate2",
327 Args: json.RawMessage(`{}`),
328 })
329
330 h.drainUntil(t, "open_panel")
331
332 h.sendToExt(t, extproto.PanelKeyFromHost{
333 Type: "panel_key", PanelID: pid, Key: "rune", Text: "n",
334 })
335
336 f := h.drainUntil(t, "tool_result")
337 var tr extproto.ToolResultFromExt
338 if err := json.Unmarshal(f.raw, &tr); err != nil {
339 t.Fatalf("unmarshal tool_result: %v", err)
340 }
341 if !tr.IsError {
342 t.Errorf("expected is_error=true on denial")
343 }
344 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "denied") {
345 t.Errorf("expected 'denied' in content, got %+v", tr.Content)
346 }
347
348 h.hostW.Close()
349}