1package extensions
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "io"
9 "os"
10 "sync"
11 "testing"
12 "testing/synctest"
13 "time"
14
15 "github.com/patriceckhart/zot/packages/agent/extproto"
16)
17
18type interactiveHooks struct{ stubHooks }
19
20func (*interactiveHooks) SupportsInteractiveTools() bool { return true }
21
22func toolPeer(t *testing.T) (*Manager, *Extension, <-chan []byte, *io.PipeWriter) {
23 t.Helper()
24 m := New(t.TempDir(), "", "test", "", "", &interactiveHooks{})
25 frames, ext := lifecyclePeer(t, m)
26 ext.pendingTool = map[string]chan extproto.ToolResultFromExt{}
27 ext.toolCancel = true
28 ext.readyCh = make(chan struct{})
29 log, err := os.CreateTemp(t.TempDir(), "log")
30 if err != nil {
31 t.Fatal(err)
32 }
33 ext.logFile = log
34 r, w := io.Pipe()
35 done := make(chan struct{})
36 go func() {
37 m.readLoop(ext, bufio.NewScanner(r))
38 close(done)
39 }()
40 t.Cleanup(func() {
41 w.Close()
42 <-done
43 r.Close()
44 log.Close()
45 })
46 m.toolIndex["ask"] = ext
47 return m, ext, frames, w
48}
49
50func sendToolFrame(t *testing.T, w io.Writer, frame any) {
51 t.Helper()
52 b, err := extproto.Encode(frame)
53 if err != nil {
54 t.Fatal(err)
55 }
56 if _, err := w.Write(b); err != nil {
57 t.Fatal(err)
58 }
59}
60
61func toolCallFrame(t *testing.T, frames <-chan []byte) extproto.ToolCallFromHost {
62 t.Helper()
63 var call extproto.ToolCallFromHost
64 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &call); err != nil {
65 t.Fatal(err)
66 }
67 if call.Type != "tool_call" {
68 t.Fatalf("expected tool_call, got %+v", call)
69 }
70 return call
71}
72
73func awaitToolError(t *testing.T, done <-chan error) error {
74 t.Helper()
75 select {
76 case err := <-done:
77 return err
78 case <-time.After(2 * time.Second):
79 t.Fatal("invocation did not finish")
80 return nil
81 }
82}
83
84func TestInteractiveToolRegistrationAndCompletion(t *testing.T) {
85 m, _, frames, replies := toolPeer(t)
86 sendToolFrame(t, replies, extproto.RegisterToolFromExt{Type: "register_tool", Name: "ask", Schema: json.RawMessage(`{}`), Interactive: true})
87 sendToolFrame(t, replies, extproto.ReadyFromExt{Type: "ready"})
88 m.WaitForReady(time.Second)
89 infos := m.Tools()
90 if len(infos) != 1 || !infos[0].Interactive {
91 t.Fatalf("registration: %+v", infos)
92 }
93 tool := NewTool(m, infos[0]).(*extensionTool)
94 if tool.timeout != 0 {
95 t.Fatal("interactive timeout is enabled")
96 }
97 if normal := NewTool(m, ToolInfo{}).(*extensionTool); normal.timeout != 60*time.Second {
98 t.Fatal("normal timeout changed")
99 }
100 done := make(chan error, 1)
101 go func() {
102 result, err := tool.Execute(context.Background(), nil, nil)
103 if err == nil && result.IsError {
104 err = errors.New("interactive completion failed")
105 }
106 done <- err
107 }()
108 call := toolCallFrame(t, frames)
109 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
110 if err := awaitToolError(t, done); err != nil {
111 t.Fatal(err)
112 }
113}
114
115func TestInteractiveToolTermination(t *testing.T) {
116 for _, reason := range []string{"cancel", "timeout", "disconnect", "shutdown"} {
117 t.Run(reason, func(t *testing.T) {
118 synctest.Test(t, func(t *testing.T) {
119 m, ext, frames, replies := toolPeer(t)
120 ctx, cancel := context.WithCancel(context.Background())
121 defer cancel()
122 timeout := time.Duration(0)
123 if reason == "timeout" {
124 timeout = time.Nanosecond
125 }
126 done := make(chan error, 1)
127 go func() {
128 _, err := m.InvokeTool(ctx, "ask", nil, timeout)
129 done <- err
130 }()
131 call := toolCallFrame(t, frames)
132 // Ensure request delivery has completed before cancelling the reply wait.
133 synctest.Wait()
134 switch reason {
135 case "cancel":
136 cancel()
137 case "disconnect":
138 replies.Close()
139 case "shutdown":
140 m.Stop(time.Second)
141 }
142 err := awaitToolError(t, done)
143 if err == nil {
144 t.Fatal("expected terminal error")
145 }
146 if reason == "cancel" && !errors.Is(err, context.Canceled) {
147 t.Fatal(err)
148 }
149 if reason == "timeout" && !errors.Is(err, context.DeadlineExceeded) {
150 t.Fatal(err)
151 }
152 ext.mu.Lock()
153 pending := len(ext.pendingTool)
154 ext.mu.Unlock()
155 if pending != 0 {
156 t.Fatal("pending call leaked")
157 }
158 if reason == "cancel" || reason == "timeout" {
159 var frame extproto.ToolCancelFromHost
160 json.Unmarshal(nextLifecycleFrame(t, frames), &frame)
161 if frame.Type != "tool_cancel" || frame.ID != call.ID {
162 t.Fatalf("cancel: %+v", frame)
163 }
164 // A late result must not block delivery of the next result.
165 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
166 go func() {
167 _, err := m.InvokeTool(context.Background(), "ask", nil, 0)
168 done <- err
169 }()
170 next := toolCallFrame(t, frames)
171 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: next.ID})
172 if err := awaitToolError(t, done); err != nil {
173 t.Fatal(err)
174 }
175 }
176 })
177 })
178 }
179}
180
181func TestInteractiveToolConcurrentReplies(t *testing.T) {
182 m, _, frames, replies := toolPeer(t)
183 const count = 32
184 done := make(chan error, count)
185 for n := 0; n < count; n++ {
186 go func() {
187 _, err := m.InvokeTool(context.Background(), "ask", nil, 0)
188 done <- err
189 }()
190 }
191 ids := make([]string, count)
192 seen := map[string]bool{}
193 for n := range ids {
194 ids[n] = toolCallFrame(t, frames).ID
195 if seen[ids[n]] {
196 t.Fatal("duplicate invocation ID")
197 }
198 seen[ids[n]] = true
199 }
200 for n := count - 1; n >= 0; n-- {
201 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: ids[n]})
202 }
203 for range ids {
204 if err := awaitToolError(t, done); err != nil {
205 t.Fatal(err)
206 }
207 }
208}
209
210func TestInteractiveToolWaitsBeyondDefaultTimeout(t *testing.T) {
211 synctest.Test(t, func(t *testing.T) {
212 m, ext, frames, replies := toolPeer(t)
213 done := make(chan error, 1)
214 go func() {
215 _, err := m.InvokeTool(context.Background(), "ask", nil, 0)
216 done <- err
217 }()
218 call := toolCallFrame(t, frames)
219 synctest.Wait()
220 // Advance the fake clock, not wall-clock time.
221 time.Sleep(2 * time.Minute)
222 synctest.Wait()
223 select {
224 case err := <-done:
225 t.Fatalf("interactive call ended early: %v", err)
226 default:
227 }
228 ext.mu.Lock()
229 pending := len(ext.pendingTool)
230 ext.mu.Unlock()
231 if pending != 1 {
232 t.Fatal("interactive call is no longer pending")
233 }
234 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID})
235 if err := awaitToolError(t, done); err != nil {
236 t.Fatal(err)
237 }
238 })
239}
240
241func TestInteractiveToolHeadless(t *testing.T) {
242 m := New(t.TempDir(), "", "test", "", "", &stubHooks{})
243 tool := NewTool(m, ToolInfo{Name: "ask", Interactive: true})
244 result, err := tool.Execute(context.Background(), nil, nil)
245 if err != nil || !result.IsError || result.Status != "failed" {
246 t.Fatalf("headless result: %+v, %v", result, err)
247 }
248}
249
250type blockedToolWriter struct {
251 started chan struct{}
252 closed chan struct{}
253 once sync.Once
254}
255
256func (w *blockedToolWriter) Write([]byte) (int, error) {
257 close(w.started)
258 <-w.closed
259 return 0, io.ErrClosedPipe
260}
261
262func (w *blockedToolWriter) Close() error {
263 w.once.Do(func() { close(w.closed) })
264 return nil
265}
266
267func TestToolCancellationDuringWrite(t *testing.T) {
268 w := &blockedToolWriter{started: make(chan struct{}), closed: make(chan struct{})}
269 pipe := newOrderedPipe(w)
270 defer pipe.Close()
271 ctx, cancel := context.WithCancel(context.Background())
272 defer cancel()
273 done := make(chan error, 1)
274 go func() {
275 _, err := pipe.writeContext(ctx, []byte("{}\n"), time.Hour)
276 done <- err
277 }()
278 select {
279 case <-w.started:
280 case <-time.After(time.Second):
281 t.Fatal("write did not start")
282 }
283 cancel()
284 if err := awaitToolError(t, done); !errors.Is(err, context.Canceled) {
285 t.Fatalf("write cancellation: %v", err)
286 }
287}