1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "os"
8 "path/filepath"
9 "strings"
10 "testing"
11 "testing/synctest"
12 "time"
13
14 "github.com/patriceckhart/zot/packages/agent/extproto"
15)
16
17func TestToolCancelCapabilityHandshake(t *testing.T) {
18 for _, tc := range []struct {
19 name string
20 caps []string
21 want bool
22 }{
23 {name: "omitted"},
24 {name: "legacy", caps: []string{"tools", "panels"}},
25 {name: "unknown", caps: []string{"tool_cancel_future", "TOOL_CANCEL"}},
26 {name: "supported", caps: []string{"tools", "tool_cancel"}, want: true},
27 } {
28 t.Run(tc.name, func(t *testing.T) {
29 root := t.TempDir()
30 writeMockToolExtension(t, filepath.Join(root, "extensions"))
31 path := filepath.Join(root, "extensions", "tool-mock", "run.sh")
32 script, err := os.ReadFile(path)
33 if err != nil {
34 t.Fatal(err)
35 }
36 hello, err := json.Marshal(extproto.HelloFromExt{Type: "hello", Name: "tool-mock", Version: "0.1", Capabilities: tc.caps})
37 if err != nil {
38 t.Fatal(err)
39 }
40 lines := strings.Split(string(script), "\n")
41 lines[1] = "printf '%s\\n' '" + string(hello) + "'"
42 if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0o755); err != nil {
43 t.Fatal(err)
44 }
45 m := New(root, "", "test", "", "", &stubHooks{})
46 t.Cleanup(func() { m.Stop(time.Second) })
47 if errs := m.Discover(context.Background()); len(errs) != 0 {
48 t.Fatal(errs)
49 }
50 m.WaitForReady(time.Second)
51 m.mu.RLock()
52 ext := m.toolIndex["echo"]
53 m.mu.RUnlock()
54 if ext == nil {
55 t.Fatal("tool was not registered")
56 }
57 if ext.toolCancel != tc.want {
58 t.Fatalf("cancellation support = %v, want %v", ext.toolCancel, tc.want)
59 }
60 })
61 }
62}
63
64func TestToolCancelCapabilityDelivery(t *testing.T) {
65 for _, supported := range []bool{false, true} {
66 name := "legacy"
67 if supported {
68 name = "supported"
69 }
70 for _, outcome := range []string{"cancel", "timeout", "complete"} {
71 t.Run(name+"/"+outcome, func(t *testing.T) {
72 synctest.Test(t, func(t *testing.T) {
73 m, ext, frames, replies := toolPeer(t)
74 ext.toolCancel = supported
75 ctx, cancel := context.WithCancel(context.Background())
76 defer cancel()
77 timeout := time.Duration(0)
78 if outcome == "timeout" {
79 timeout = time.Nanosecond
80 }
81 done := make(chan error, 1)
82 go func() {
83 _, err := m.InvokeTool(ctx, "ask", nil, timeout)
84 done <- err
85 }()
86 call := toolCallFrame(t, frames)
87 synctest.Wait()
88 var want error
89 switch outcome {
90 case "cancel":
91 want = context.Canceled
92 cancel()
93 case "timeout":
94 want = context.DeadlineExceeded
95 case "complete":
96 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
97 }
98 if err := awaitToolError(t, done); !errors.Is(err, want) {
99 t.Fatalf("invocation error = %v, want %v", err, want)
100 }
101 if supported && outcome != "complete" {
102 var frame extproto.ToolCancelFromHost
103 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &frame); err != nil {
104 t.Fatal(err)
105 }
106 if frame.Type != "tool_cancel" || frame.ID != call.ID {
107 t.Fatalf("unexpected cancellation frame: %+v", frame)
108 }
109 }
110 synctest.Wait()
111 select {
112 case frame := <-frames:
113 t.Fatalf("unexpected host frame: %s", frame)
114 default:
115 }
116 ext.mu.Lock()
117 pending := len(ext.pendingTool)
118 ext.mu.Unlock()
119 if pending != 0 {
120 t.Fatal("pending invocation leaked")
121 }
122 // Abandoned results remain harmless without cancellation support.
123 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
124 go func() {
125 _, err := m.InvokeTool(context.Background(), "ask", nil, 0)
126 done <- err
127 }()
128 next := toolCallFrame(t, frames)
129 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
130 if err := awaitToolError(t, done); err != nil {
131 t.Fatal(err)
132 }
133 })
134 })
135 }
136 }
137}