1package extensions
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "io"
9 "os"
10 "strings"
11 "sync"
12 "testing"
13 "time"
14
15 "github.com/bnema/zut/packages/agent/extproto"
16 "github.com/bnema/zut/packages/core"
17 "github.com/bnema/zut/packages/provider"
18)
19
20// interactiveHooks marks the host as able to show interactive tool prompts.
21type interactiveHooks struct{ stubHooks }
22
23func (*interactiveHooks) SupportsInteractiveTools() bool { return true }
24
25// toolPeer wires a Manager to an in-process fake extension over io.Pipe
26// pairs. The test plays the extension: it reads host frames from `frames`
27// and replies through `replies`. The ordered transport is real, so write
28// bounds, disconnects, and cancel frames exercise production code.
29type toolPeer struct {
30 m *Manager
31 ext *Extension
32 frames chan []byte // host -> extension frames, one per element
33 replies *io.PipeWriter // extension -> host (feeds readLoop)
34 readDone chan struct{}
35}
36
37func newToolPeer(t *testing.T, hooks HostHooks, toolCancel bool) *toolPeer {
38 return newToolPeerOpts(t, hooks, toolCancel, true)
39}
40
41// newToolPeerOpts builds the peer; drain=false leaves the host's stdin pipe
42// unread so every host write blocks in the transport writer.
43func newToolPeerOpts(t *testing.T, hooks HostHooks, toolCancel, drain bool) *toolPeer {
44 t.Helper()
45 m := New(t.TempDir(), "", "0.0.0-test", "", "", hooks)
46 log, err := os.CreateTemp(t.TempDir(), "log")
47 if err != nil {
48 t.Fatal(err)
49 }
50 // Host stdin side: manager writes, test reads.
51 stdinR, stdinW := io.Pipe()
52 // Host stdout side: test writes, manager readLoop reads.
53 stdoutR, stdoutW := io.Pipe()
54 ext := &Extension{
55 Manifest: Manifest{Name: "peer"},
56 readyCh: make(chan struct{}),
57 pending: map[string]chan extproto.CommandResponseFromExt{},
58 pendingTool: map[string]chan extproto.ToolResultFromExt{},
59 pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{},
60 eventSubs: map[string]struct{}{},
61 interceptSubs: map[string]struct{}{},
62 lifecycleFrames: make(chan []byte, 64),
63 lifecycleStop: make(chan struct{}),
64 logFile: log,
65 helloAck: true,
66 toolCancel: toolCancel,
67 }
68 ext.stdin = newOrderedPipe(stdinW)
69 m.mu.Lock()
70 m.ext[ext.Manifest.Name] = ext
71 m.mu.Unlock()
72
73 p := &toolPeer{m: m, ext: ext, frames: make(chan []byte, 64), replies: stdoutW, readDone: make(chan struct{})}
74 go func() {
75 m.readLoop(ext, bufio.NewScanner(stdoutR))
76 close(p.readDone)
77 }()
78 if drain {
79 go func() {
80 sc := bufio.NewScanner(stdinR)
81 sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
82 for sc.Scan() {
83 p.frames <- append([]byte(nil), sc.Bytes()...)
84 }
85 close(p.frames)
86 }()
87 }
88 t.Cleanup(func() {
89 _ = stdoutW.Close()
90 <-p.readDone
91 _ = stdinR.Close()
92 _ = log.Close()
93 })
94 return p
95}
96
97func (p *toolPeer) send(t *testing.T, frame any) {
98 t.Helper()
99 b, err := extproto.Encode(frame)
100 if err != nil {
101 t.Fatal(err)
102 }
103 if _, err := p.replies.Write(b); err != nil {
104 t.Fatal(err)
105 }
106}
107
108func (p *toolPeer) register(t *testing.T, name string, interactive bool) {
109 t.Helper()
110 p.send(t, extproto.RegisterToolFromExt{Type: "register_tool", Name: name, Schema: json.RawMessage(`{"type":"object"}`), Interactive: interactive})
111 p.send(t, extproto.ReadyFromExt{Type: "ready"})
112 p.m.WaitForReady(2 * time.Second)
113}
114
115func (p *toolPeer) next(t *testing.T) extproto.Frame {
116 t.Helper()
117 select {
118 case raw, ok := <-p.frames:
119 if !ok {
120 t.Fatal("host transport closed before expected frame")
121 }
122 var f extproto.Frame
123 if err := json.Unmarshal(raw, &f); err != nil {
124 t.Fatal(err)
125 }
126 return f
127 case <-time.After(2 * time.Second):
128 t.Fatal("no frame from host")
129 return extproto.Frame{}
130 }
131}
132
133func (p *toolPeer) expectToolCall(t *testing.T) extproto.Frame {
134 t.Helper()
135 f := p.next(t)
136 if f.Type != "tool_call" {
137 t.Fatalf("expected tool_call, got %+v", f)
138 }
139 return f
140}
141
142// expectClosedWithout drains until the host transport closes and fails if a
143// frame of the given type was seen. Used to prove legacy peers never see
144// tool_cancel and headless hosts never dispatch interactive calls.
145func (p *toolPeer) expectClosedWithout(t *testing.T, frameType string) {
146 t.Helper()
147 for {
148 select {
149 case raw, ok := <-p.frames:
150 if !ok {
151 return
152 }
153 var f extproto.Frame
154 _ = json.Unmarshal(raw, &f)
155 if f.Type == frameType {
156 t.Fatalf("host sent %s: %s", frameType, raw)
157 }
158 case <-time.After(2 * time.Second):
159 t.Fatal("host transport did not close")
160 }
161 }
162}
163
164func execTool(tool core.Tool, ctx context.Context) <-chan struct {
165 res core.ToolResult
166 err error
167} {
168 done := make(chan struct {
169 res core.ToolResult
170 err error
171 }, 1)
172 go func() {
173 res, err := tool.Execute(ctx, nil, nil)
174 done <- struct {
175 res core.ToolResult
176 err error
177 }{res, err}
178 }()
179 return done
180}
181
182func resultText(res core.ToolResult) string {
183 var b strings.Builder
184 for _, c := range res.Content {
185 if tb, ok := c.(provider.TextBlock); ok {
186 b.WriteString(tb.Text)
187 }
188 }
189 return b.String()
190}
191
192func TestInteractiveToolRegistrationAndTimeoutPolicy(t *testing.T) {
193 p := newToolPeer(t, &interactiveHooks{}, true)
194 p.send(t, extproto.RegisterToolFromExt{Type: "register_tool", Name: "ask", Schema: json.RawMessage(`{}`), Interactive: true})
195 p.send(t, extproto.RegisterToolFromExt{Type: "register_tool", Name: "plain", Schema: json.RawMessage(`{}`)})
196 p.send(t, extproto.ReadyFromExt{Type: "ready"})
197 p.m.WaitForReady(2 * time.Second)
198
199 infos := p.m.Tools()
200 if len(infos) != 2 {
201 t.Fatalf("tools = %+v", infos)
202 }
203 for _, info := range infos {
204 tool := NewTool(p.m, info).(*extensionTool)
205 switch info.Name {
206 case "ask":
207 if !info.Interactive || tool.timeout != 0 || !tool.Interactive() {
208 t.Fatalf("interactive tool policy wrong: info=%+v timeout=%v", info, tool.timeout)
209 }
210 case "plain":
211 if info.Interactive || tool.timeout != 60*time.Second || tool.Interactive() {
212 t.Fatalf("normal tool policy wrong: info=%+v timeout=%v", info, tool.timeout)
213 }
214 }
215 }
216}
217
218func TestInteractiveToolCompletesWithoutDeadline(t *testing.T) {
219 p := newToolPeer(t, &interactiveHooks{}, true)
220 p.register(t, "ask", true)
221 tool := NewTool(p.m, p.m.Tools()[0])
222
223 done := execTool(tool, context.Background())
224 call := p.expectToolCall(t)
225 p.send(t, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "answer"}}})
226 out := <-done
227 if out.err != nil || out.res.IsError || resultText(out.res) != "answer" {
228 t.Fatalf("unexpected result: %+v err=%v", out.res, out.err)
229 }
230 p.ext.mu.Lock()
231 pending := len(p.ext.pendingTool)
232 p.ext.mu.Unlock()
233 if pending != 0 {
234 t.Fatalf("pending slots leaked: %d", pending)
235 }
236}
237
238func TestInteractiveToolResultSurvivesImmediateDisconnect(t *testing.T) {
239 for range 30 {
240 p := newToolPeer(t, &interactiveHooks{}, true)
241 p.register(t, "ask", true)
242 done := execTool(NewTool(p.m, p.m.Tools()[0]), context.Background())
243 call := p.expectToolCall(t)
244 p.send(t, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "answer"}}})
245 if err := p.replies.Close(); err != nil {
246 t.Fatal(err)
247 }
248 <-p.readDone
249 select {
250 case out := <-done:
251 if out.err != nil || out.res.IsError || resultText(out.res) != "answer" {
252 t.Fatalf("completed reply lost to disconnect: result=%+v err=%v", out.res, out.err)
253 }
254 case <-time.After(2 * time.Second):
255 t.Fatal("tool call remained blocked after disconnect")
256 }
257 }
258}
259
260func TestInteractiveToolContextCancelSendsToolCancelAndIgnoresLateResult(t *testing.T) {
261 p := newToolPeer(t, &interactiveHooks{}, true)
262 p.register(t, "ask", true)
263 tool := NewTool(p.m, p.m.Tools()[0])
264
265 ctx, cancel := context.WithCancel(context.Background())
266 done := execTool(tool, ctx)
267 call := p.expectToolCall(t)
268 cancel()
269 out := <-done
270 if out.err != nil || !out.res.IsError || !strings.Contains(resultText(out.res), "cancelled") {
271 t.Fatalf("expected cancelled tool error, got %+v err=%v", out.res, out.err)
272 }
273 f := p.next(t)
274 if f.Type != "tool_cancel" || f.ID != call.ID {
275 t.Fatalf("expected tool_cancel for %s, got %+v", call.ID, f)
276 }
277 // A late reply for the abandoned id must be ignored and the pending map
278 // must stay empty.
279 p.send(t, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
280 // Prove the read loop is still healthy by completing a fresh call.
281 done = execTool(tool, context.Background())
282 call2 := p.expectToolCall(t)
283 if call2.ID == call.ID {
284 t.Fatal("correlation ids repeated")
285 }
286 p.send(t, extproto.ToolResultFromExt{Type: "tool_result", ID: call2.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "ok"}}})
287 if out := <-done; out.err != nil || out.res.IsError {
288 t.Fatalf("follow-up call failed: %+v err=%v", out.res, out.err)
289 }
290 p.ext.mu.Lock()
291 pending := len(p.ext.pendingTool)
292 p.ext.mu.Unlock()
293 if pending != 0 {
294 t.Fatalf("pending slots leaked: %d", pending)
295 }
296}
297
298func TestInteractiveToolContextDeadlineEndsCall(t *testing.T) {
299 p := newToolPeer(t, &interactiveHooks{}, true)
300 p.register(t, "ask", true)
301 tool := NewTool(p.m, p.m.Tools()[0])
302
303 ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
304 defer cancel()
305 done := execTool(tool, ctx)
306 call := p.expectToolCall(t)
307 out := <-done
308 if out.err != nil || !out.res.IsError || !strings.Contains(resultText(out.res), "deadline") {
309 t.Fatalf("expected deadline tool error, got %+v err=%v", out.res, out.err)
310 }
311 if f := p.next(t); f.Type != "tool_cancel" || f.ID != call.ID {
312 t.Fatalf("expected tool_cancel, got %+v", f)
313 }
314}
315
316func TestNormalToolReplyTimeoutSendsToolCancel(t *testing.T) {
317 p := newToolPeer(t, &interactiveHooks{}, true)
318 p.register(t, "plain", false)
319
320 done := make(chan error, 1)
321 go func() {
322 _, err := p.m.InvokeTool(context.Background(), "plain", json.RawMessage(`{}`), 20*time.Millisecond)
323 done <- err
324 }()
325 call := p.expectToolCall(t)
326 err := <-done
327 if !errors.Is(err, context.DeadlineExceeded) || !strings.Contains(err.Error(), "timeout waiting for peer/plain") {
328 t.Fatalf("expected reply timeout, got %v", err)
329 }
330 if f := p.next(t); f.Type != "tool_cancel" || f.ID != call.ID {
331 t.Fatalf("expected tool_cancel, got %+v", f)
332 }
333 // The wrapper reports a plain timeout as an ordinary tool error.
334 tool := NewTool(p.m, p.m.Tools()[0]).(*extensionTool)
335 tool.timeout = 20 * time.Millisecond
336 res, execErr := tool.Execute(context.Background(), nil, nil)
337 p.expectToolCall(t)
338 if execErr != nil || !res.IsError || !strings.Contains(resultText(res), "failed: timeout") {
339 t.Fatalf("expected timeout tool error, got %+v err=%v", res, execErr)
340 }
341}
342
343func TestLegacyPeerNeverReceivesToolCancel(t *testing.T) {
344 p := newToolPeer(t, &interactiveHooks{}, false)
345 p.register(t, "ask", true)
346 tool := NewTool(p.m, p.m.Tools()[0])
347
348 ctx, cancel := context.WithCancel(context.Background())
349 done := execTool(tool, ctx)
350 p.expectToolCall(t)
351 cancel()
352 if out := <-done; out.err != nil || !out.res.IsError {
353 t.Fatalf("expected cancelled tool error, got %+v err=%v", out.res, out.err)
354 }
355 // Legacy registration without the interactive field keeps the old
356 // timeout, and cancellation never produces a tool_cancel frame.
357 p.m.Stop(time.Second)
358 p.expectClosedWithout(t, "tool_cancel")
359}
360
361func TestInteractiveToolDisconnectEndsCall(t *testing.T) {
362 p := newToolPeer(t, &interactiveHooks{}, true)
363 p.register(t, "ask", true)
364 tool := NewTool(p.m, p.m.Tools()[0])
365
366 done := execTool(tool, context.Background())
367 p.expectToolCall(t)
368 // Extension exits: its stdout closes, readLoop ends and tears the
369 // transport down.
370 _ = p.replies.Close()
371 out := <-done
372 if out.err != nil || !out.res.IsError || !strings.Contains(resultText(out.res), "disconnected") {
373 t.Fatalf("expected disconnect tool error, got %+v err=%v", out.res, out.err)
374 }
375 if p.m.HasTool("ask") {
376 t.Fatal("tool still registered after disconnect")
377 }
378}
379
380func TestInteractiveToolShutdownEndsCall(t *testing.T) {
381 p := newToolPeer(t, &interactiveHooks{}, true)
382 p.register(t, "ask", true)
383 tool := NewTool(p.m, p.m.Tools()[0])
384
385 done := execTool(tool, context.Background())
386 p.expectToolCall(t)
387 stopped := make(chan struct{})
388 go func() {
389 p.m.Stop(200 * time.Millisecond)
390 close(stopped)
391 }()
392 out := <-done
393 if out.err != nil || !out.res.IsError {
394 t.Fatalf("expected tool error on shutdown, got %+v err=%v", out.res, out.err)
395 }
396 select {
397 case <-stopped:
398 case <-time.After(3 * time.Second):
399 t.Fatal("Stop blocked on the pending interactive call")
400 }
401}
402
403func TestInteractiveToolBlockedWriteIsBoundedAndCancellable(t *testing.T) {
404 // The extension never reads stdin, so the tool_call write blocks in the
405 // transport writer. Context cancellation must return promptly instead of
406 // waiting on the (deadline-free) reply, and must not pin goroutines.
407 t.Run("cancel", func(t *testing.T) {
408 p := newToolPeerOpts(t, &interactiveHooks{}, true, false)
409 p.register(t, "ask", true)
410 tool := NewTool(p.m, p.m.Tools()[0])
411 ctx, cancel := context.WithCancel(context.Background())
412 done := execTool(tool, ctx)
413 // Wait for the frame to reach the writer goroutine (blocked on the
414 // unread pipe), then cancel.
415 pipe := p.ext.stdin.(*orderedPipe)
416 for {
417 pipe.mu.Lock()
418 queued := len(pipe.queue)
419 pipe.mu.Unlock()
420 if queued == 0 {
421 break
422 }
423 time.Sleep(time.Millisecond)
424 }
425 cancel()
426 select {
427 case out := <-done:
428 if out.err != nil || !out.res.IsError {
429 t.Fatalf("expected tool error, got %+v err=%v", out.res, out.err)
430 }
431 case <-time.After(3 * time.Second):
432 t.Fatal("cancelled call stayed blocked on the transport write")
433 }
434 })
435 // Without cancellation the write itself is still bounded, independent of
436 // the absent reply deadline: the transport times out and disconnects.
437 t.Run("timeout", func(t *testing.T) {
438 p := newToolPeerOpts(t, &interactiveHooks{}, true, false)
439 p.register(t, "ask", true)
440 pipe := p.ext.stdin.(*orderedPipe)
441 done := make(chan error, 1)
442 go func() {
443 _, err := pipe.writeContext(context.Background(), []byte("{}\n"), 30*time.Millisecond)
444 done <- err
445 }()
446 select {
447 case err := <-done:
448 if !errors.Is(err, context.DeadlineExceeded) {
449 t.Fatalf("expected transport timeout, got %v", err)
450 }
451 case <-time.After(3 * time.Second):
452 t.Fatal("blocked write not bounded")
453 }
454 // A subsequent tool call fails immediately on the dead transport.
455 tool := NewTool(p.m, p.m.Tools()[0])
456 res, err := tool.Execute(context.Background(), nil, nil)
457 if err != nil || !res.IsError || !strings.Contains(resultText(res), "write:") {
458 t.Fatalf("expected write failure, got %+v err=%v", res, err)
459 }
460 })
461}
462
463func TestOrderedPipeWriteTimeoutDisconnects(t *testing.T) {
464 r, w := io.Pipe()
465 pipe := newOrderedPipe(w)
466 t.Cleanup(func() { _ = r.Close() })
467 // Nobody reads r, so the first write blocks in the writer goroutine and
468 // the second waits in the queue.
469 first := make(chan error, 1)
470 go func() {
471 _, err := pipe.writeContext(context.Background(), []byte("a\n"), 30*time.Millisecond)
472 first <- err
473 }()
474 select {
475 case err := <-first:
476 if !errors.Is(err, context.DeadlineExceeded) {
477 t.Fatalf("expected timeout, got %v", err)
478 }
479 case <-time.After(3 * time.Second):
480 t.Fatal("bounded write did not return")
481 }
482 select {
483 case <-pipe.Done():
484 case <-time.After(time.Second):
485 t.Fatal("transport not disconnected after write timeout")
486 }
487 if _, err := pipe.Write([]byte("b\n")); !errors.Is(err, io.ErrClosedPipe) {
488 t.Fatalf("write after disconnect = %v", err)
489 }
490}
491
492func TestOrderedPipeCancelWithdrawsQueuedFrameWithoutDisconnect(t *testing.T) {
493 r, w := io.Pipe()
494 pipe := newOrderedPipe(w)
495 t.Cleanup(func() { pipe.Close(); _ = r.Close() })
496
497 // Head frame blocks in the writer; the second frame stays queued.
498 headAck := make(chan error, 1)
499 if _, err := pipe.enqueue([]byte("head\n"), headAck); err != nil {
500 t.Fatal(err)
501 }
502 ctx, cancel := context.WithCancel(context.Background())
503 second := make(chan error, 1)
504 go func() {
505 _, err := pipe.writeContext(ctx, []byte("second\n"), time.Minute)
506 second <- err
507 }()
508 cancel()
509 select {
510 case err := <-second:
511 if !errors.Is(err, context.Canceled) {
512 t.Fatalf("expected cancel, got %v", err)
513 }
514 case <-time.After(3 * time.Second):
515 t.Fatal("queued write did not return on cancel")
516 }
517 select {
518 case <-pipe.Done():
519 t.Fatal("cancel of a queued frame must not disconnect the transport")
520 default:
521 }
522 // Reader drains: only the head frame is delivered, the withdrawn one is
523 // not, and the transport keeps working.
524 go func() {
525 _, _ = pipe.Write([]byte("third\n"))
526 }()
527 sc := bufio.NewScanner(r)
528 var got []string
529 for len(got) < 2 && sc.Scan() {
530 got = append(got, sc.Text())
531 }
532 if strings.Join(got, ",") != "head,third" {
533 t.Fatalf("delivered frames = %v", got)
534 }
535}
536
537func TestOrderedPipeOverflowDisconnects(t *testing.T) {
538 r, w := io.Pipe()
539 pipe := newOrderedPipe(w)
540 t.Cleanup(func() { _ = r.Close() })
541 var err error
542 for i := 0; i <= outboundFrames+1; i++ {
543 if _, err = pipe.enqueue([]byte("x\n"), nil); err != nil {
544 break
545 }
546 }
547 if !errors.Is(err, errOutboundOverflow) {
548 t.Fatalf("expected overflow, got %v", err)
549 }
550 select {
551 case <-pipe.Done():
552 case <-time.After(time.Second):
553 t.Fatal("overflow did not disconnect")
554 }
555}
556
557func TestCorrelationIDsAreUniqueUnderConcurrency(t *testing.T) {
558 const n = 2000
559 var mu sync.Mutex
560 seen := make(map[string]struct{}, n)
561 var wg sync.WaitGroup
562 for i := 0; i < n; i++ {
563 wg.Add(1)
564 go func() {
565 defer wg.Done()
566 id := newCorrelationID()
567 mu.Lock()
568 seen[id] = struct{}{}
569 mu.Unlock()
570 }()
571 }
572 wg.Wait()
573 if len(seen) != n {
574 t.Fatalf("got %d unique ids for %d calls", len(seen), n)
575 }
576}
577
578func TestHeadlessHostsRejectInteractiveTools(t *testing.T) {
579 cases := map[string]HostHooks{
580 "missing_method": &stubHooks{},
581 "returns_false": &interactiveOptOut{},
582 }
583 for name, hooks := range cases {
584 t.Run(name, func(t *testing.T) {
585 p := newToolPeer(t, hooks, true)
586 p.register(t, "ask", true)
587 tool := NewTool(p.m, p.m.Tools()[0])
588 res, err := tool.Execute(context.Background(), nil, nil)
589 if err != nil || !res.IsError || !strings.Contains(resultText(res), "requires an interactive host") {
590 t.Fatalf("expected fail-closed rejection, got %+v err=%v", res, err)
591 }
592 // Nothing was sent to the extension.
593 p.m.Stop(time.Second)
594 p.expectClosedWithout(t, "tool_call")
595 })
596 }
597 // A non-interactive tool still works on a headless host.
598 p := newToolPeer(t, &stubHooks{}, true)
599 p.register(t, "plain", false)
600 tool := NewTool(p.m, p.m.Tools()[0])
601 done := execTool(tool, context.Background())
602 call := p.expectToolCall(t)
603 p.send(t, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "ok"}}})
604 if out := <-done; out.err != nil || out.res.IsError {
605 t.Fatalf("normal tool on headless host failed: %+v err=%v", out.res, out.err)
606 }
607}
608
609type interactiveOptOut struct{ stubHooks }
610
611func (*interactiveOptOut) SupportsInteractiveTools() bool { return false }
612
613func TestIsInteractiveToolClassifiesExtensionTools(t *testing.T) {
614 m := New(t.TempDir(), "", "0.0.0-test", "", "", &stubHooks{})
615 if !IsInteractiveTool(NewTool(m, ToolInfo{Name: "ask", Interactive: true})) {
616 t.Fatal("interactive tool not classified")
617 }
618 if IsInteractiveTool(NewTool(m, ToolInfo{Name: "plain"})) {
619 t.Fatal("plain tool classified as interactive")
620 }
621 if IsInteractiveTool(nil) {
622 t.Fatal("nil classified as interactive")
623 }
624}