1package extensions
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "io"
8 "reflect"
9 "testing"
10 "time"
11
12 "github.com/patriceckhart/zot/packages/agent/extproto"
13)
14
15func lifecyclePeer(t *testing.T, m *Manager, names ...string) (<-chan []byte, *Extension) {
16 t.Helper()
17 reader, writer := io.Pipe()
18 out := make(chan []byte, 512)
19 done := make(chan struct{})
20 pipe := newOrderedPipe(writer)
21 ext := &Extension{stdin: pipe, eventSubs: map[string]struct{}{}, interceptSubs: map[string]struct{}{}, pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{}}
22 for _, name := range names {
23 ext.eventSubs[name] = struct{}{}
24 }
25 m.ext["test"] = ext
26 go func() {
27 defer close(done)
28 scanner := bufio.NewScanner(reader)
29 for scanner.Scan() {
30 out <- append([]byte(nil), scanner.Bytes()...)
31 }
32 close(out)
33 }()
34 t.Cleanup(func() { pipe.Close(); reader.Close(); <-done })
35 return out, ext
36}
37
38func nextLifecycleFrame(t *testing.T, frames <-chan []byte) []byte {
39 t.Helper()
40 select {
41 case frame, ok := <-frames:
42 if !ok {
43 t.Fatal("peer closed")
44 }
45 return frame
46 case <-time.After(2 * time.Second):
47 t.Fatal("timed out waiting for frame")
48 return nil
49 }
50}
51
52func TestLifecycleOrderAndShutdownDrain(t *testing.T) {
53 m := New(t.TempDir(), "/work", "test", "", "", nil)
54 frames, _ := lifecyclePeer(t, m, "session_start", "session_end", "tool_call", "tool_result")
55 m.StartSession("one", "/work", "")
56 m.StartSession("one", "/work", "") // idempotent
57 for n := 0; n < 50; n++ {
58 m.EmitEvent(extproto.EventFromHost{Event: "tool_call", Step: n})
59 m.EmitEvent(extproto.EventFromHost{Event: "tool_result", Step: n})
60 }
61 m.StartSession("two", "/other", "cwd_change")
62 m.EndSession("user_exit")
63 m.EndSession("shutdown")
64 m.Stop(time.Second)
65 m.StartSession("late", "/late", "")
66 m.EmitEvent(extproto.EventFromHost{Event: "tool_result"})
67 if id, _ := m.SessionContext(); id != "two" {
68 t.Fatal("late callback reopened a stopped session")
69 }
70 var got []extproto.EventFromHost
71 for {
72 raw := nextLifecycleFrame(t, frames)
73 var e extproto.EventFromHost
74 if err := json.Unmarshal(raw, &e); err != nil {
75 t.Fatal(err)
76 }
77 if e.Type == "shutdown" {
78 break
79 }
80 if e.Sequence != uint64(len(got)+1) {
81 t.Fatalf("out of order: %+v", e)
82 }
83 got = append(got, e)
84 }
85 if len(got) != 104 {
86 t.Fatalf("got %d events", len(got))
87 }
88 for n := 0; n < 50; n++ {
89 a, b := got[1+2*n], got[2+2*n]
90 if a.Event != "tool_call" || b.Event != "tool_result" || a.Step != n || b.Step != n {
91 t.Fatalf("wrong order at %d", n)
92 }
93 }
94 if got[101].Event != "session_end" || got[101].SessionID != "one" || got[101].Reason != "cwd_change" || got[101].CWD != "/work" {
95 t.Fatalf("old boundary: %+v", got[101])
96 }
97 if got[102].Event != "session_start" || got[102].SessionID != "two" || got[102].CWD != "/other" {
98 t.Fatalf("new boundary: %+v", got[102])
99 }
100 if got[103].Event != "session_end" || got[103].Reason != "user_exit" {
101 t.Fatalf("end: %+v", got[103])
102 }
103}
104
105func TestNotificationPrecedesInterceptOnWire(t *testing.T) {
106 m := New(t.TempDir(), "/work", "test", "", "", nil)
107 frames, ext := lifecyclePeer(t, m, "tool_call")
108 ext.interceptSubs["tool_call"] = struct{}{}
109 m.EmitEvent(extproto.EventFromHost{Event: "tool_call", ToolID: "call"})
110 result := make(chan InterceptResult, 1)
111 go func() { result <- m.InterceptToolCall(context.Background(), "call", "bash", json.RawMessage(`{}`)) }()
112 var event extproto.EventFromHost
113 json.Unmarshal(nextLifecycleFrame(t, frames), &event)
114 if event.Type != "event" || event.Event != "tool_call" {
115 t.Fatalf("first frame: %+v", event)
116 }
117 var intercept extproto.EventInterceptFromHost
118 json.Unmarshal(nextLifecycleFrame(t, frames), &intercept)
119 if intercept.Type != "event_intercept" {
120 t.Fatalf("second frame: %+v", intercept)
121 }
122 ext.mu.Lock()
123 reply := ext.pendingIntercept[intercept.ID]
124 delete(ext.pendingIntercept, intercept.ID)
125 ext.mu.Unlock()
126 reply <- extproto.EventInterceptResponseFromExt{Block: true}
127 select {
128 case r := <-result:
129 if !r.Block {
130 t.Fatal("block lost")
131 }
132 case <-time.After(time.Second):
133 t.Fatal("interceptor did not finish")
134 }
135}
136
137func TestSlowConsumerOverflowClosesPipe(t *testing.T) {
138 reader, writer := io.Pipe()
139 defer reader.Close()
140 pipe := newOrderedPipe(writer)
141 defer pipe.Close()
142 var err error
143 for n := 0; n < outboundFrames+2; n++ {
144 err = pipe.enqueue([]byte("{}\n"), nil)
145 if err != nil {
146 break
147 }
148 }
149 if err == nil {
150 t.Fatal("unbounded queue")
151 }
152 select {
153 case <-pipe.done:
154 default:
155 t.Fatal("overflow did not disconnect")
156 }
157 if _, err := pipe.Write([]byte("{}\n")); err == nil {
158 t.Fatal("write after disconnect succeeded")
159 }
160}
161
162func TestLifecycleSubscriptionFiltering(t *testing.T) {
163 m := New(t.TempDir(), "/work", "test", "", "", nil)
164 frames, ext := lifecyclePeer(t, m, "post_compact", "permission_decision", "subagent_start", "subagent_stop", "user_prompt_submit")
165 names := []string{"post_compact", "permission_decision", "subagent_start", "subagent_stop", "user_prompt_submit"}
166 m.EmitEvent(extproto.EventFromHost{Event: "pre_compact"})
167 for _, name := range names {
168 m.EmitEvent(extproto.EventFromHost{Event: name})
169 }
170 if _, err := ext.stdin.Write([]byte("{\"type\":\"marker\"}\n")); err != nil {
171 t.Fatal(err)
172 }
173 var got []string
174 for range names {
175 var e extproto.EventFromHost
176 json.Unmarshal(nextLifecycleFrame(t, frames), &e)
177 got = append(got, e.Event)
178 }
179 if !reflect.DeepEqual(got, names) {
180 t.Fatalf("got %v", got)
181 }
182}
183
184func TestMalformedArgumentsStillEmitResult(t *testing.T) {
185 m := New(t.TempDir(), "/work", "test", "", "", nil)
186 frames, _ := lifecyclePeer(t, m, "tool_result")
187 m.EmitEvent(extproto.EventFromHost{Event: "tool_result", ToolID: "bad-call", ToolArgs: json.RawMessage(`{invalid`), Status: "failed"})
188 var e extproto.EventFromHost
189 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &e); err != nil {
190 t.Fatal(err)
191 }
192 if e.ToolID != "bad-call" || e.Status != "failed" || len(e.ToolArgs) != 0 || e.ToolArgsRaw != `{invalid` {
193 t.Fatalf("result: %+v", e)
194 }
195}
196
197func TestShutdownBoundsBlockedWriter(t *testing.T) {
198 reader, writer := io.Pipe()
199 defer reader.Close()
200 pipe := newOrderedPipe(writer)
201 defer pipe.Close()
202 if err := pipe.enqueue([]byte("{}\n"), nil); err != nil {
203 t.Fatal(err)
204 }
205 done := make(chan struct{})
206 go func() { stopExtensions([]*Extension{{stdin: pipe}}, time.Millisecond); close(done) }()
207 select {
208 case <-done:
209 case <-time.After(time.Second):
210 t.Fatal("shutdown blocked on a slow consumer")
211 }
212}
213
214func TestExtensionToolTerminalStatus(t *testing.T) {
215 for _, cancelled := range []bool{false, true} {
216 name := "timed_out"
217 if cancelled {
218 name = "cancelled"
219 }
220 t.Run(name, func(t *testing.T) {
221 m := New(t.TempDir(), "/work", "test", "", "", nil)
222 _, ext := lifecyclePeer(t, m)
223 ext.pendingTool = map[string]chan extproto.ToolResultFromExt{}
224 m.toolIndex["probe"] = ext
225 ctx, cancel := context.WithCancel(context.Background())
226 defer cancel()
227 timeout := time.Nanosecond
228 if cancelled {
229 cancel()
230 timeout = time.Hour
231 }
232 tool := &extensionTool{name: "probe", extension: "test", manager: m, timeout: timeout}
233 result, err := tool.Execute(ctx, json.RawMessage(`{}`), nil)
234 if err != nil || !result.IsError || result.Status != name {
235 t.Fatalf("result: %+v, error: %v", result, err)
236 }
237 ext.mu.Lock()
238 pending := len(ext.pendingTool)
239 ext.mu.Unlock()
240 if pending != 0 {
241 t.Fatal("terminal call left a pending reply")
242 }
243 })
244 }
245}