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// TestOpenPanelEmitsCorrectFrame checks that e.OpenPanel sends a
141// well-formed open_panel frame with the correct PanelSpec fields.
142func TestOpenPanelEmitsCorrectFrame(t *testing.T) {
143 h := newHarness("test-ext")
144 go h.ext.Run()
145 h.handshake(t)
146
147 go h.ext.OpenPanel("my-panel", "My Title", []string{"line a", "line b"}, "esc close")
148
149 f := h.drainUntil(t, "open_panel")
150
151 var op extproto.OpenPanelFromExt
152 if err := json.Unmarshal(f.raw, &op); err != nil {
153 t.Fatalf("unmarshal open_panel: %v", err)
154 }
155 if op.Panel.ID != "my-panel" {
156 t.Errorf("panel id: want %q, got %q", "my-panel", op.Panel.ID)
157 }
158 if op.Panel.Title != "My Title" {
159 t.Errorf("panel title: want %q, got %q", "My Title", op.Panel.Title)
160 }
161 if len(op.Panel.Lines) != 2 || op.Panel.Lines[0] != "line a" || op.Panel.Lines[1] != "line b" {
162 t.Errorf("panel lines: got %v", op.Panel.Lines)
163 }
164 if op.Panel.Footer != "esc close" {
165 t.Errorf("panel footer: want %q, got %q", "esc close", op.Panel.Footer)
166 }
167
168 h.hostW.Close()
169}
170
171// TestBlockingToolWaitsForPanelKey is the core integration test for
172// the human-in-the-loop pattern: the tool handler opens a panel,
173// blocks on a channel, and only returns a tool_result after a key
174// event arrives.
175func TestBlockingToolWaitsForPanelKey(t *testing.T) {
176 h := newHarness("gate-ext")
177
178 const pid = "gate-panel"
179 const toolCallID = "tc-001"
180
181 approved := make(chan bool, 1)
182
183 h.ext.OnPanelKey(pid, func(key, text string) {
184 switch {
185 case key == "rune" && text == "y":
186 h.ext.ClosePanel(pid)
187 approved <- true
188 case key == "rune" && text == "n", key == "esc":
189 h.ext.ClosePanel(pid)
190 approved <- false
191 }
192 }, func() { approved <- false })
193
194 h.ext.Tool("gate", "needs approval",
195 json.RawMessage(`{"type":"object","properties":{}}`),
196 func(args json.RawMessage) ToolResult {
197 h.ext.OpenPanel(pid, "Approve?",
198 []string{" y approve", " n deny"}, "y/n")
199 if <-approved {
200 return TextResult("approved")
201 }
202 return TextErrorResult("denied")
203 })
204
205 go h.ext.Run()
206 h.handshake(t)
207
208 h.sendToExt(t, extproto.ToolCallFromHost{
209 Type: "tool_call", ID: toolCallID, Name: "gate",
210 Args: json.RawMessage(`{}`),
211 })
212
213 // Tool goroutine must open the panel before it can reply.
214 h.drainUntil(t, "open_panel")
215
216 // Send approval — tool should now unblock and emit tool_result.
217 h.sendToExt(t, extproto.PanelKeyFromHost{
218 Type: "panel_key", PanelID: pid, Key: "rune", Text: "y",
219 })
220
221 f := h.drainUntil(t, "tool_result")
222 var tr extproto.ToolResultFromExt
223 if err := json.Unmarshal(f.raw, &tr); err != nil {
224 t.Fatalf("unmarshal tool_result: %v", err)
225 }
226 if tr.ID != toolCallID {
227 t.Errorf("tool_result id: want %q, got %q", toolCallID, tr.ID)
228 }
229 if tr.IsError {
230 t.Errorf("expected success, got is_error=true")
231 }
232 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "approved") {
233 t.Errorf("expected 'approved' in content, got %+v", tr.Content)
234 }
235
236 h.hostW.Close()
237}
238
239// TestBlockingToolDenied mirrors TestBlockingToolWaitsForPanelKey but
240// sends "n" so the tool returns an error result.
241func TestBlockingToolDenied(t *testing.T) {
242 h := newHarness("gate-ext-deny")
243
244 const pid = "deny-panel"
245 const toolCallID = "tc-002"
246
247 approved := make(chan bool, 1)
248
249 h.ext.OnPanelKey(pid, func(key, text string) {
250 switch {
251 case key == "rune" && text == "y":
252 h.ext.ClosePanel(pid)
253 approved <- true
254 case key == "rune" && text == "n", key == "esc":
255 h.ext.ClosePanel(pid)
256 approved <- false
257 }
258 }, func() { approved <- false })
259
260 h.ext.Tool("gate2", "needs approval",
261 json.RawMessage(`{"type":"object","properties":{}}`),
262 func(args json.RawMessage) ToolResult {
263 h.ext.OpenPanel(pid, "Approve?", []string{"y/n"}, "")
264 if <-approved {
265 return TextResult("approved")
266 }
267 return TextErrorResult("denied")
268 })
269
270 go h.ext.Run()
271 h.handshake(t)
272
273 h.sendToExt(t, extproto.ToolCallFromHost{
274 Type: "tool_call", ID: toolCallID, Name: "gate2",
275 Args: json.RawMessage(`{}`),
276 })
277
278 h.drainUntil(t, "open_panel")
279
280 h.sendToExt(t, extproto.PanelKeyFromHost{
281 Type: "panel_key", PanelID: pid, Key: "rune", Text: "n",
282 })
283
284 f := h.drainUntil(t, "tool_result")
285 var tr extproto.ToolResultFromExt
286 if err := json.Unmarshal(f.raw, &tr); err != nil {
287 t.Fatalf("unmarshal tool_result: %v", err)
288 }
289 if !tr.IsError {
290 t.Errorf("expected is_error=true on denial")
291 }
292 if len(tr.Content) == 0 || !strings.Contains(tr.Content[0].Text, "denied") {
293 t.Errorf("expected 'denied' in content, got %+v", tr.Content)
294 }
295
296 h.hostW.Close()
297}