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 defer m.TrackToolCall("call", "caller")()
110 m.EmitEvent(extproto.EventFromHost{Event: "tool_call", ToolID: "call", OriginExtension: "caller"})
111 result := make(chan InterceptResult, 1)
112 go func() { result <- m.InterceptToolCall(context.Background(), "call", "bash", json.RawMessage(`{}`)) }()
113 var event extproto.EventFromHost
114 json.Unmarshal(nextLifecycleFrame(t, frames), &event)
115 if event.Type != "event" || event.Event != "tool_call" {
116 t.Fatalf("first frame: %+v", event)
117 }
118 var intercept extproto.EventInterceptFromHost
119 json.Unmarshal(nextLifecycleFrame(t, frames), &intercept)
120 if intercept.Type != "event_intercept" || intercept.OriginExtension != "caller" {
121 t.Fatalf("second frame: %+v", intercept)
122 }
123 ext.mu.Lock()
124 reply := ext.pendingIntercept[intercept.ID]
125 delete(ext.pendingIntercept, intercept.ID)
126 ext.mu.Unlock()
127 reply <- extproto.EventInterceptResponseFromExt{Block: true}
128 select {
129 case r := <-result:
130 if !r.Block {
131 t.Fatal("block lost")
132 }
133 case <-time.After(time.Second):
134 t.Fatal("interceptor did not finish")
135 }
136}
137
138func TestSlowConsumerOverflowClosesPipe(t *testing.T) {
139 reader, writer := io.Pipe()
140 defer reader.Close()
141 pipe := newOrderedPipe(writer)
142 defer pipe.Close()
143 var err error
144 for n := 0; n < outboundFrames+2; n++ {
145 err = pipe.enqueue([]byte("{}\n"), nil)
146 if err != nil {
147 break
148 }
149 }
150 if err == nil {
151 t.Fatal("unbounded queue")
152 }
153 select {
154 case <-pipe.done:
155 default:
156 t.Fatal("overflow did not disconnect")
157 }
158 if _, err := pipe.Write([]byte("{}\n")); err == nil {
159 t.Fatal("write after disconnect succeeded")
160 }
161}
162
163func TestLifecycleSubscriptionFiltering(t *testing.T) {
164 m := New(t.TempDir(), "/work", "test", "", "", nil)
165 frames, ext := lifecyclePeer(t, m, "post_compact", "permission_decision", "subagent_start", "subagent_stop", "user_prompt_submit")
166 names := []string{"post_compact", "permission_decision", "subagent_start", "subagent_stop", "user_prompt_submit"}
167 m.EmitEvent(extproto.EventFromHost{Event: "pre_compact"})
168 for _, name := range names {
169 m.EmitEvent(extproto.EventFromHost{Event: name})
170 }
171 if _, err := ext.stdin.Write([]byte("{\"type\":\"marker\"}\n")); err != nil {
172 t.Fatal(err)
173 }
174 var got []string
175 for range names {
176 var e extproto.EventFromHost
177 json.Unmarshal(nextLifecycleFrame(t, frames), &e)
178 got = append(got, e.Event)
179 }
180 if !reflect.DeepEqual(got, names) {
181 t.Fatalf("got %v", got)
182 }
183}
184
185func TestMalformedArgumentsStillEmitResult(t *testing.T) {
186 m := New(t.TempDir(), "/work", "test", "", "", nil)
187 frames, _ := lifecyclePeer(t, m, "tool_result")
188 m.EmitEvent(extproto.EventFromHost{Event: "tool_result", ToolID: "bad-call", ToolArgs: json.RawMessage(`{invalid`), Status: "failed"})
189 var e extproto.EventFromHost
190 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &e); err != nil {
191 t.Fatal(err)
192 }
193 if e.ToolID != "bad-call" || e.Status != "failed" || len(e.ToolArgs) != 0 || e.ToolArgsRaw != `{invalid` {
194 t.Fatalf("result: %+v", e)
195 }
196}
197
198func TestShutdownBoundsBlockedWriter(t *testing.T) {
199 reader, writer := io.Pipe()
200 defer reader.Close()
201 pipe := newOrderedPipe(writer)
202 defer pipe.Close()
203 if err := pipe.enqueue([]byte("{}\n"), nil); err != nil {
204 t.Fatal(err)
205 }
206 done := make(chan struct{})
207 go func() { stopExtensions([]*Extension{{stdin: pipe}}, time.Millisecond); close(done) }()
208 select {
209 case <-done:
210 case <-time.After(time.Second):
211 t.Fatal("shutdown blocked on a slow consumer")
212 }
213}
214
215func TestExtensionToolTerminalStatus(t *testing.T) {
216 for _, cancelled := range []bool{false, true} {
217 name := "timed_out"
218 if cancelled {
219 name = "cancelled"
220 }
221 t.Run(name, func(t *testing.T) {
222 m := New(t.TempDir(), "/work", "test", "", "", nil)
223 _, ext := lifecyclePeer(t, m)
224 ext.pendingTool = map[string]chan extproto.ToolResultFromExt{}
225 m.toolIndex["probe"] = ext
226 ctx, cancel := context.WithCancel(context.Background())
227 defer cancel()
228 timeout := time.Nanosecond
229 if cancelled {
230 cancel()
231 timeout = time.Hour
232 }
233 tool := &extensionTool{name: "probe", extension: "test", manager: m, timeout: timeout}
234 result, err := tool.Execute(ctx, json.RawMessage(`{}`), nil)
235 if err != nil || !result.IsError || result.Status != name {
236 t.Fatalf("result: %+v, error: %v", result, err)
237 }
238 ext.mu.Lock()
239 pending := len(ext.pendingTool)
240 ext.mu.Unlock()
241 if pending != 0 {
242 t.Fatal("terminal call left a pending reply")
243 }
244 })
245 }
246}