1package extensions
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "io"
8 "os"
9 "strings"
10 "testing"
11 "time"
12
13 "github.com/patriceckhart/zot/packages/agent/extproto"
14)
15
16type startWriter struct{ write func([]byte) (int, error) }
17
18func (w startWriter) Write(p []byte) (int, error) { return w.write(p) }
19func (w startWriter) Close() error { return nil }
20
21func addStartExtension(t *testing.T, m *Manager, name string, response func(extproto.EventInterceptFromHost) string) *Extension {
22 t.Helper()
23 log, err := os.CreateTemp(t.TempDir(), "log")
24 if err != nil {
25 t.Fatal(err)
26 }
27 reader, writer := io.Pipe()
28 ext := &Extension{Manifest: Manifest{Name: name}, logFile: log, readyCh: make(chan struct{}), pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{}, interceptSubs: map[string]struct{}{}, eventSubs: map[string]struct{}{}}
29 ext.stdin = startWriter{write: func(p []byte) (int, error) {
30 var ev extproto.EventInterceptFromHost
31 if err := json.Unmarshal(p, &ev); err != nil {
32 return 0, err
33 }
34 _, err := io.WriteString(writer, response(ev)+"\n")
35 return len(p), err
36 }}
37 m.mu.Lock()
38 m.ext[name] = ext
39 m.extOrder = append(m.extOrder, ext)
40 m.mu.Unlock()
41 done := make(chan struct{})
42 scanner := bufio.NewScanner(reader)
43 scanner.Buffer(make([]byte, 4096), 4*1024*1024)
44 go func() { defer close(done); m.readLoop(ext, scanner) }()
45 t.Cleanup(func() { writer.Close(); <-done; reader.Close(); log.Close() })
46 if _, err := io.WriteString(writer, `{"type":"subscribe","intercept":["before_agent_start"]}`+"\n"+`{"type":"ready"}`+"\n"); err != nil {
47 t.Fatal(err)
48 }
49 <-ext.readyCh
50 return ext
51}
52
53func TestBeforeAgentStartResponses(t *testing.T) {
54 for _, tc := range []struct {
55 name, raw, want string
56 warning bool
57 }{
58 {"unchanged", `"base"`, "base", false},
59 {"appended", `"base plus"`, "base plus", false},
60 {"rewritten", `"BASE"`, "BASE", false},
61 {"replaced", `" völlig neu\n\n"`, " völlig neu\n\n", false},
62 {"empty", `""`, "", false},
63 {"omitted", "", "base", false},
64 {"null", "null", "base", true},
65 {"number", "42", "base", true},
66 {"object", "{}", "base", true},
67 {"oversized", `"` + strings.Repeat("x", MaxSystemPromptResponseSize) + `"`, "base", true},
68 } {
69 t.Run(tc.name, func(t *testing.T) {
70 hooks := &stubHooks{}
71 m := New(t.TempDir(), "/effective", "test", "provider", "model", hooks)
72 ext := addStartExtension(t, m, "test", func(ev extproto.EventInterceptFromHost) string {
73 if ev.Event != "before_agent_start" || ev.SystemPrompt == nil || *ev.SystemPrompt != "base" || ev.CWD != "/effective" || ev.SessionID != "session" || ev.AgentRunID != "run" || ev.Provider != "provider" || ev.Model != "model" {
74 t.Errorf("wrong event metadata")
75 }
76 field := ""
77 if tc.raw != "" {
78 field = `,"system_prompt":` + tc.raw
79 }
80 return `{"type":"event_intercept_response","id":"` + ev.ID + `"` + field + `}`
81 })
82 base := "base"
83 got := m.InterceptBeforeAgentStart(context.Background(), extproto.EventInterceptFromHost{SystemPrompt: &base, SessionID: "session", AgentRunID: "run", Model: "model"})
84 if got != tc.want {
85 t.Fatalf("got %q, want %q", got, tc.want)
86 }
87 ext.mu.Lock()
88 warnings := len(ext.diagnostics)
89 ext.mu.Unlock()
90 if (warnings > 0) != tc.warning {
91 t.Fatalf("warnings = %d", warnings)
92 }
93 hooks.mu.Lock()
94 notifications := len(hooks.notifies)
95 hooks.mu.Unlock()
96 if (notifications > 0) != tc.warning {
97 t.Fatalf("notifications = %d", notifications)
98 }
99 })
100 }
101}
102
103func TestBeforeAgentStartChainsInLoadOrder(t *testing.T) {
104 m := New(t.TempDir(), "", "test", "p", "m", nil)
105 for _, name := range []string{"z-first", "a-second"} {
106 addStartExtension(t, m, name, func(ev extproto.EventInterceptFromHost) string {
107 raw, _ := json.Marshal(*ev.SystemPrompt + "/" + name)
108 return `{"type":"event_intercept_response","id":"` + ev.ID + `","system_prompt":` + string(raw) + `}`
109 })
110 }
111 for range 20 {
112 base := "base"
113 if got := m.InterceptBeforeAgentStart(context.Background(), extproto.EventInterceptFromHost{SystemPrompt: &base}); got != "base/z-first/a-second" {
114 t.Fatal(got)
115 }
116 }
117}
118
119func TestBeforeAgentStartMalformedResponseTimesOut(t *testing.T) {
120 hooks := &stubHooks{}
121 m := New(t.TempDir(), "", "test", "p", "m", hooks)
122 ext := addStartExtension(t, m, "bad", func(extproto.EventInterceptFromHost) string { return `{invalid` })
123 addStartExtension(t, m, "good", func(ev extproto.EventInterceptFromHost) string {
124 if *ev.SystemPrompt != "base" {
125 t.Error("failed extension changed prompt")
126 }
127 return `{"type":"event_intercept_response","id":"` + ev.ID + `","system_prompt":"good"}`
128 })
129 ctx, cancel := context.WithTimeout(context.Background(), 2*interceptTimeout)
130 defer cancel()
131 base := "base"
132 if got := m.InterceptBeforeAgentStart(ctx, extproto.EventInterceptFromHost{SystemPrompt: &base}); got != "good" {
133 t.Fatalf("chain did not continue after timeout: %q", got)
134 }
135 ext.mu.Lock()
136 pending := len(ext.pendingIntercept)
137 ext.mu.Unlock()
138 if pending != 0 {
139 t.Fatal("pending request leaked")
140 }
141 hooks.mu.Lock()
142 defer hooks.mu.Unlock()
143 if len(hooks.notifies) != 1 || !strings.Contains(hooks.notifies[0], "timed out") {
144 t.Fatalf("missing timeout warning: %v", hooks.notifies)
145 }
146}
147
148func TestBeforeAgentStartBoundsBlockedWrite(t *testing.T) {
149 m := New(t.TempDir(), "", "test", "p", "m", nil)
150 log, err := os.CreateTemp(t.TempDir(), "log")
151 if err != nil {
152 t.Fatal(err)
153 }
154 defer log.Close()
155 reader, writer := io.Pipe()
156 defer reader.Close()
157 ext := &Extension{stdin: writer, logFile: log, pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{}}
158 ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
159 defer cancel()
160 result := m.askIntercept(ctx, ext, extproto.EventInterceptFromHost{Event: "before_agent_start"})
161 if result.failure == "" {
162 t.Fatal("expected write timeout")
163 }
164 if _, err := writer.Write([]byte("x")); err == nil {
165 t.Fatal("blocked pipe not closed")
166 }
167}