1package agent
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "os"
8 "path/filepath"
9 "strings"
10 "testing"
11 "time"
12 "unicode/utf8"
13
14 "github.com/patriceckhart/zot/packages/agent/extensions"
15 "github.com/patriceckhart/zot/packages/agent/extproto"
16 "github.com/patriceckhart/zot/packages/core"
17 "github.com/patriceckhart/zot/packages/provider"
18)
19
20// This subprocess is a real protocol peer. All data is synthetic and lives in
21// t.TempDir; the parent reads the captured wire frames only after shutdown.
22func TestLifecycleProcessHelper(t *testing.T) {
23 if os.Getenv("ZOT_TEST_LIFECYCLE") != "1" {
24 return
25 }
26 file, err := os.Create(os.Getenv("ZOT_TEST_LIFECYCLE_OUTPUT"))
27 if err != nil {
28 os.Exit(2)
29 }
30 enc := json.NewEncoder(os.Stdout)
31 enc.Encode(extproto.HelloFromExt{Type: "hello", Name: "lifecycle", Version: "1"})
32 enc.Encode(extproto.SubscribeFromExt{Type: "subscribe", Events: []string{"session_start", "session_end", "user_prompt_submit", "turn_start", "turn_end", "tool_call", "tool_result", "permission_decision", "pre_compact", "post_compact"}, Intercept: []string{"tool_call"}})
33 enc.Encode(extproto.ReadyFromExt{Type: "ready"})
34 scanner := bufio.NewScanner(os.Stdin)
35 for scanner.Scan() {
36 var e extproto.EventInterceptFromHost
37 if json.Unmarshal(scanner.Bytes(), &e) != nil {
38 os.Exit(3)
39 }
40 switch e.Type {
41 case "event":
42 if _, err := file.Write(append(append([]byte(nil), scanner.Bytes()...), '\n')); err != nil {
43 os.Exit(4)
44 }
45 case "event_intercept":
46 enc.Encode(extproto.EventInterceptResponseFromExt{Type: "event_intercept_response", ID: e.ID, Block: os.Getenv("ZOT_TEST_LIFECYCLE_BLOCK") == "1", Reason: "test policy", ModifiedArgs: json.RawMessage(`{"effective":true}`)})
47 case "shutdown":
48 file.Close()
49 enc.Encode(extproto.ShutdownAckFromExt{Type: "shutdown_ack"})
50 os.Exit(0)
51 }
52 }
53 file.Close()
54 os.Exit(0)
55}
56
57type lifecycleClient struct{ step int }
58
59func (*lifecycleClient) Name() string { return "test" }
60func (c *lifecycleClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
61 c.step++
62 stop := provider.StopEnd
63 content := []provider.Content{provider.TextBlock{Text: "done"}}
64 if c.step == 1 {
65 stop = provider.StopToolUse
66 content = []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "probe", Arguments: json.RawMessage(`{"original":true}`)}}
67 }
68 ch := make(chan provider.Event, 1)
69 ch <- provider.EventDone{Stop: stop, Message: provider.Message{Role: provider.RoleAssistant, Content: content}}
70 close(ch)
71 return ch, nil
72}
73
74type lifecycleTool struct{ args json.RawMessage }
75
76func (*lifecycleTool) Name() string { return "probe" }
77func (*lifecycleTool) Description() string { return "test" }
78func (*lifecycleTool) Schema() json.RawMessage { return json.RawMessage(`{}`) }
79func (t *lifecycleTool) Execute(_ context.Context, args json.RawMessage, _ func(string)) (core.ToolResult, error) {
80 t.args = append(json.RawMessage(nil), args...)
81 return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: "ok"}}}, nil
82}
83
84func TestExtensionLifecycleEndToEnd(t *testing.T) {
85 for _, blocked := range []bool{false, true} {
86 name := "rewritten"
87 if blocked {
88 name = "blocked"
89 }
90 t.Run(name, func(t *testing.T) {
91 dir := t.TempDir()
92 output := filepath.Join(dir, "events.jsonl")
93 t.Setenv("ZOT_TEST_LIFECYCLE", "1")
94 t.Setenv("ZOT_TEST_LIFECYCLE_OUTPUT", output)
95 if blocked {
96 t.Setenv("ZOT_TEST_LIFECYCLE_BLOCK", "1")
97 } else {
98 t.Setenv("ZOT_TEST_LIFECYCLE_BLOCK", "0")
99 }
100 exe, err := os.Executable()
101 if err != nil {
102 t.Fatal(err)
103 }
104 manifest, err := json.Marshal(extensions.Manifest{Name: "lifecycle", Exec: exe, Args: []string{"-test.run=^TestLifecycleProcessHelper$"}})
105 if err != nil {
106 t.Fatal(err)
107 }
108 if err := os.WriteFile(filepath.Join(dir, "extension.json"), manifest, 0600); err != nil {
109 t.Fatal(err)
110 }
111 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
112 defer cancel()
113 mgr := extensions.New(t.TempDir(), dir, "test", "test", "test", nil)
114 if errs := mgr.LoadExplicit(ctx, []string{dir}); len(errs) > 0 {
115 t.Fatal(errs)
116 }
117 t.Cleanup(func() { mgr.Stop(time.Second) })
118 mgr.WaitForReady(time.Second)
119 tool := &lifecycleTool{}
120 ag := core.NewAgent(&lifecycleClient{}, "test", "", core.NewRegistry(tool))
121 ag.SessionID = "session-one"
122 wireNonInteractiveAgentExtHooks(ctx, ag, mgr)
123 if err := ag.Prompt(ctx, "review", nil, func(e core.AgentEvent) {
124 if _, ok := e.(core.EvPromptSubmit); ok {
125 t.Error("extension-only submission leaked into the RPC/SDK stream")
126 }
127 }); err != nil {
128 t.Fatal(err)
129 }
130 // Exercise the same explicit boundary helper used by the session picker.
131 ag.SessionID = "session-two"
132 startExtensionSession(mgr, ag, dir, "session_switch")
133 mgr.Stop(time.Second)
134 data, err := os.ReadFile(output)
135 if err != nil {
136 t.Fatal(err)
137 }
138 var events []extproto.EventFromHost
139 scanner := bufio.NewScanner(strings.NewReader(string(data)))
140 for scanner.Scan() {
141 var e extproto.EventFromHost
142 if err := json.Unmarshal(scanner.Bytes(), &e); err != nil {
143 t.Fatal(err)
144 }
145 if e.Sequence == 0 || (len(events) > 0 && e.Sequence <= events[len(events)-1].Sequence) {
146 t.Fatalf("wire sequence: %+v", e)
147 }
148 events = append(events, e)
149 }
150 if err := scanner.Err(); err != nil {
151 t.Fatal(err)
152 }
153 indexes := map[string]int{}
154 for n, e := range events {
155 if _, exists := indexes[e.Event]; !exists {
156 indexes[e.Event] = n
157 }
158 if e.SessionID == "" || e.CWD == "" {
159 t.Fatalf("missing correlation: %+v", e)
160 }
161 }
162 names := []string{"session_start", "user_prompt_submit", "turn_start", "tool_call", "turn_end", "permission_decision", "tool_result", "session_end"}
163 previous := -1
164 for _, name := range names {
165 n, ok := indexes[name]
166 if !ok || n <= previous {
167 t.Fatalf("missing/out-of-order %s: %+v", name, events)
168 }
169 previous = n
170 }
171 result := events[indexes["tool_result"]]
172 if result.ToolID != "call-1" || result.ToolName != "probe" || result.Result == nil || result.Executed == nil {
173 t.Fatalf("result: %+v", result)
174 }
175 if blocked {
176 if result.Status != "blocked" || *result.Executed || tool.args != nil || !result.Result.IsError {
177 t.Fatalf("blocked outcome: %+v", result)
178 }
179 } else {
180 if result.Status != "completed" || !*result.Executed || string(result.ToolArgs) != string(tool.args) || string(tool.args) != `{"effective":true}` {
181 t.Fatalf("effective result: %+v", result)
182 }
183 }
184 tail := events[len(events)-3:]
185 if tail[0].Reason != "session_switch" || tail[0].SessionID != "session-one" || tail[1].Event != "session_start" || tail[1].SessionID != "session-two" || tail[2].Reason != "shutdown" {
186 t.Fatalf("session boundaries: %+v", tail)
187 }
188 })
189 }
190}
191
192func TestEventResultPayloadBudget(t *testing.T) {
193 result := core.ToolResult{Content: []provider.Content{
194 provider.TextBlock{Text: strings.Repeat("ü", eventResultBudget)},
195 provider.ImageBlock{MimeType: "image/png", Data: []byte("synthetic")},
196 }}
197 out := extensionEventResult(result)
198 if !out.Truncated || len(out.Content) != 1 || len(out.Content[0].Text) > eventResultBudget || !utf8.ValidString(out.Content[0].Text) {
199 t.Fatalf("invalid bounded result")
200 }
201 if len(result.Content[0].(provider.TextBlock).Text) != 2*eventResultBudget {
202 t.Fatal("transcript modified")
203 }
204 small := extensionEventResult(core.ToolResult{Content: []provider.Content{provider.ImageBlock{MimeType: "image/png", Data: []byte("test")}}})
205 if small.Truncated || len(small.Content) != 1 || small.Content[0].Data != "dGVzdA==" {
206 t.Fatalf("image: %+v", small)
207 }
208}