1package ext
2
3import (
4 "context"
5 "encoding/json"
6 "testing"
7 "time"
8
9 "github.com/bnema/zut/packages/agent/extproto"
10)
11
12// interactiveHarness registers one interactive tool whose handler blocks on
13// ctx or on an explicit answer, and reports when it started and ended.
14type interactiveHarness struct {
15 *extHarness
16 started chan string // tool_call id
17 ended chan error // ctx.Err() at handler exit
18 answers chan string
19}
20
21func newInteractiveHarness(t *testing.T) *interactiveHarness {
22 t.Helper()
23 h := &interactiveHarness{
24 extHarness: newHarness("interactive-ext"),
25 started: make(chan string, 8),
26 ended: make(chan error, 8),
27 answers: make(chan string, 8),
28 }
29 h.ext.InteractiveTool("ask", "ask the user", json.RawMessage(`{"type":"object"}`),
30 func(ctx context.Context, args json.RawMessage) ToolResult {
31 var in struct {
32 ID string `json:"id"`
33 }
34 _ = json.Unmarshal(args, &in)
35 h.started <- in.ID
36 select {
37 case a := <-h.answers:
38 h.ended <- ctx.Err()
39 return TextResult(a)
40 case <-ctx.Done():
41 h.ended <- ctx.Err()
42 return TextErrorResult("cancelled")
43 }
44 })
45 h.ext.Tool("plain", "plain", json.RawMessage(`{"type":"object"}`), func(json.RawMessage) ToolResult {
46 return TextResult("plain-ok")
47 })
48 return h
49}
50
51func (h *interactiveHarness) call(t *testing.T, id string) {
52 t.Helper()
53 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: id, Name: "ask", Args: json.RawMessage(`{"id":"` + id + `"}`)})
54 select {
55 case got := <-h.started:
56 if got != id {
57 t.Fatalf("handler started for %q, want %q", got, id)
58 }
59 case <-time.After(2 * time.Second):
60 t.Fatalf("handler for %s did not start", id)
61 }
62}
63
64func (h *interactiveHarness) awaitEnd(t *testing.T) error {
65 t.Helper()
66 select {
67 case err := <-h.ended:
68 return err
69 case <-time.After(2 * time.Second):
70 t.Fatal("handler did not end")
71 return nil
72 }
73}
74
75// expectNoToolResult drains frames for a short window and fails on any
76// tool_result. Used after a cancel to prove late results are suppressed.
77func (h *interactiveHarness) expectNoToolResult(t *testing.T) {
78 t.Helper()
79 // Force a round trip: an event frame is processed on the read loop after
80 // any tool_result the handler goroutine could have queued, so once we
81 // observe the ack for it, an earlier tool_result would already be here.
82 h.ext.OnPanelKey("probe", func(string, string) { h.ext.Notify("info", "probe") }, nil)
83 h.sendToExt(t, extproto.PanelKeyFromHost{Type: "panel_key", PanelID: "probe", Key: "rune", Text: "x"})
84 for {
85 f := h.next(t)
86 switch f.hdr.Type {
87 case "tool_result":
88 t.Fatalf("late tool_result emitted after cancellation: %s", f.raw)
89 case "notify":
90 return
91 }
92 }
93}
94
95func TestInteractiveToolRegistersFlagAndCapability(t *testing.T) {
96 h := newInteractiveHarness(t)
97 h.startRun(t)
98 f := h.next(t)
99 var hello extproto.HelloFromExt
100 if err := json.Unmarshal(f.raw, &hello); err != nil {
101 t.Fatal(err)
102 }
103 hasCancel := false
104 for _, c := range hello.Capabilities {
105 if c == "tool_cancel" {
106 hasCancel = true
107 }
108 }
109 if !hasCancel {
110 t.Fatalf("hello capabilities = %v, want tool_cancel", hello.Capabilities)
111 }
112 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", ProtocolVersion: extproto.ProtocolVersion, ZutVersion: "0.0.0-test"})
113 seen := map[string]bool{}
114 for {
115 f := h.next(t)
116 if f.hdr.Type == "ready" {
117 break
118 }
119 if f.hdr.Type != "register_tool" {
120 continue
121 }
122 var rt extproto.RegisterToolFromExt
123 if err := json.Unmarshal(f.raw, &rt); err != nil {
124 t.Fatal(err)
125 }
126 seen[rt.Name] = rt.Interactive
127 }
128 if !seen["ask"] || seen["plain"] {
129 t.Fatalf("interactive flags = %v", seen)
130 }
131}
132
133func TestInteractiveToolCompletes(t *testing.T) {
134 h := newInteractiveHarness(t)
135 h.startRun(t)
136 h.handshake(t)
137 h.call(t, "c1")
138 h.answers <- "yes"
139 if err := h.awaitEnd(t); err != nil {
140 t.Fatalf("handler ctx err = %v, want nil", err)
141 }
142 f := h.drainUntil(t, "tool_result")
143 var tr extproto.ToolResultFromExt
144 if err := json.Unmarshal(f.raw, &tr); err != nil {
145 t.Fatal(err)
146 }
147 if tr.ID != "c1" || tr.IsError || len(tr.Content) != 1 || tr.Content[0].Text != "yes" {
148 t.Fatalf("tool_result = %+v", tr)
149 }
150 // The goroutine's deferred cleanup runs after respondTool; poll briefly.
151 deadline := time.Now().Add(2 * time.Second)
152 for {
153 h.ext.mu.Lock()
154 inflight := len(h.ext.toolCancels)
155 h.ext.mu.Unlock()
156 if inflight == 0 {
157 break
158 }
159 if time.Now().After(deadline) {
160 t.Fatalf("in-flight cancel handles leaked: %d", inflight)
161 }
162 time.Sleep(time.Millisecond)
163 }
164}
165
166func TestToolCancelCancelsHandlerAndSuppressesResult(t *testing.T) {
167 h := newInteractiveHarness(t)
168 h.startRun(t)
169 h.handshake(t)
170 h.call(t, "c1")
171 h.call(t, "c2")
172 h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "c1"})
173 if err := h.awaitEnd(t); err != context.Canceled {
174 t.Fatalf("cancelled handler ctx err = %v", err)
175 }
176 h.expectNoToolResult(t)
177 // The other in-flight call is unaffected and still completes.
178 h.answers <- "two"
179 if err := h.awaitEnd(t); err != nil {
180 t.Fatalf("second handler ctx err = %v", err)
181 }
182 f := h.drainUntil(t, "tool_result")
183 var tr extproto.ToolResultFromExt
184 _ = json.Unmarshal(f.raw, &tr)
185 if tr.ID != "c2" || tr.Content[0].Text != "two" {
186 t.Fatalf("tool_result = %+v", tr)
187 }
188 // Unknown / already-finished ids are ignored without side effects.
189 h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "c1"})
190 h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "nope"})
191 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "c3", Name: "plain", Args: json.RawMessage(`{}`)})
192 f = h.drainUntil(t, "tool_result")
193 _ = json.Unmarshal(f.raw, &tr)
194 if tr.ID != "c3" || tr.Content[0].Text != "plain-ok" {
195 t.Fatalf("plain tool after cancels = %+v", tr)
196 }
197}
198
199func TestShutdownCancelsInteractiveHandlers(t *testing.T) {
200 h := newInteractiveHarness(t)
201 h.startRun(t)
202 h.handshake(t)
203 h.call(t, "c1")
204 h.sendToExt(t, extproto.ShutdownFromHost{Type: "shutdown"})
205 if err := h.awaitEnd(t); err != context.Canceled {
206 t.Fatalf("handler ctx err on shutdown = %v", err)
207 }
208 h.drainUntil(t, "shutdown_ack")
209}
210
211func TestHostEOFCancelsInteractiveHandlers(t *testing.T) {
212 h := newInteractiveHarness(t)
213 runDone := make(chan error, 1)
214 go func() { runDone <- h.ext.Run() }()
215 h.handshake(t)
216 h.call(t, "c1")
217 if err := h.hostW.Close(); err != nil {
218 t.Fatal(err)
219 }
220 if err := h.awaitEnd(t); err != context.Canceled {
221 t.Fatalf("handler ctx err on EOF = %v", err)
222 }
223 select {
224 case <-runDone:
225 case <-time.After(2 * time.Second):
226 t.Fatal("Run did not return after EOF")
227 }
228}
229
230func TestPlainToolHandlersRemainSourceCompatible(t *testing.T) {
231 h := newInteractiveHarness(t)
232 h.startRun(t)
233 h.handshake(t)
234 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "p1", Name: "plain", Args: json.RawMessage(`{}`)})
235 f := h.drainUntil(t, "tool_result")
236 var tr extproto.ToolResultFromExt
237 _ = json.Unmarshal(f.raw, &tr)
238 if tr.ID != "p1" || tr.IsError || tr.Content[0].Text != "plain-ok" {
239 t.Fatalf("tool_result = %+v", tr)
240 }
241}