1package modes
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "net/http"
8 "net/http/httptest"
9 "os"
10 "path/filepath"
11 "testing"
12 "time"
13
14 "github.com/patriceckhart/zot/packages/agent/ext"
15 "github.com/patriceckhart/zot/packages/agent/extensions"
16 "github.com/patriceckhart/zot/packages/core"
17 "github.com/patriceckhart/zot/packages/provider"
18)
19
20type toolPromptClient struct{ requests chan provider.Request }
21
22func (*toolPromptClient) Name() string { return "tool-prompt-test" }
23func (c *toolPromptClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
24 c.requests <- req
25 out := make(chan provider.Event, 1)
26 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "done"}}}}
27 close(out)
28 return out, nil
29}
30
31type shortcutTool struct{}
32
33func (*shortcutTool) Name() string { return "skill" }
34func (*shortcutTool) Description() string { return "skill" }
35func (*shortcutTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
36func (*shortcutTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) {
37 return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: "instructions"}}}, nil
38}
39
40// TestToolPromptExtensionProcess runs only as a child started by the manager.
41func TestToolPromptExtensionProcess(t *testing.T) {
42 if os.Getenv("ZOT_TOOL_PROMPT_CHILD") != "1" {
43 return
44 }
45 e := ext.New("shortcut", "test")
46 e.Command("commit", "commit", func(args string) ext.Response {
47 return ext.ToolPrompt("skill", map[string]string{"name": "developer:commit"}, args)
48 })
49 if err := e.Run(); err != nil {
50 t.Fatal(err)
51 }
52}
53
54func TestToolPromptExtensionCommandWithLocalGemini(t *testing.T) {
55 root := t.TempDir()
56 path, err := os.Executable()
57 if err != nil {
58 t.Fatal(err)
59 }
60 dir := filepath.Join(root, "extensions", "shortcut")
61 if err := os.MkdirAll(dir, 0o755); err != nil {
62 t.Fatal(err)
63 }
64 manifest, err := json.Marshal(map[string]any{"name": "shortcut", "version": "test", "exec": path, "args": []string{"-test.run=^TestToolPromptExtensionProcess$"}})
65 if err != nil {
66 t.Fatal(err)
67 }
68 if err := os.WriteFile(filepath.Join(dir, "extension.json"), manifest, 0o644); err != nil {
69 t.Fatal(err)
70 }
71 t.Setenv("ZOT_TOOL_PROMPT_CHILD", "1")
72 mgr := extensions.New(root, "", "test", "google", "gemini-3-flash-preview", nil)
73 if errs := mgr.Discover(context.Background()); len(errs) > 0 {
74 t.Fatal(errs)
75 }
76 t.Cleanup(func() { mgr.Stop(time.Second) })
77 mgr.WaitForReady(2 * time.Second)
78 if !mgr.HasCommand("commit") {
79 t.Fatal("extension command not registered")
80 }
81 requests := make(chan json.RawMessage, 1)
82 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
83 defer r.Body.Close()
84 var raw json.RawMessage
85 if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
86 http.Error(w, err.Error(), 400)
87 return
88 }
89 requests <- raw
90 w.Header().Set("Content-Type", "text/event-stream")
91 _, _ = fmt.Fprint(w, "data: "+`{"candidates":[{"content":{"role":"model","parts":[{"text":"done"}]},"finishReason":"STOP"}]}`+"\n\n")
92 }))
93 defer server.Close()
94 agent := core.NewAgent(provider.NewGemini("synthetic-test-key", server.URL), "gemini-3-flash-preview", "", core.Registry{"skill": &shortcutTool{}})
95 iv := NewInteractive(InteractiveConfig{Agent: agent, Extensions: mgr})
96 iv.runCtx = context.Background()
97 iv.invokeExtensionCommand(context.Background(), "commit", "and push")
98 select {
99 case raw := <-requests:
100 var wire struct {
101 Contents []struct {
102 Parts []struct {
103 Text string `json:"text"`
104 Signature string `json:"thoughtSignature"`
105 FunctionCall *struct {
106 Name string `json:"name"`
107 } `json:"functionCall"`
108 FunctionResponse *struct {
109 Response map[string]string `json:"response"`
110 } `json:"functionResponse"`
111 } `json:"parts"`
112 } `json:"contents"`
113 }
114 if err := json.Unmarshal(raw, &wire); err != nil {
115 t.Fatal(err)
116 }
117 if len(wire.Contents) != 3 || wire.Contents[0].Parts[0].Text != "and push" || wire.Contents[1].Parts[0].FunctionCall == nil || wire.Contents[1].Parts[0].FunctionCall.Name != "skill" || wire.Contents[1].Parts[0].Signature != "skip_thought_signature_validator" || wire.Contents[2].Parts[0].FunctionResponse == nil || wire.Contents[2].Parts[0].FunctionResponse.Response["output"] != "instructions" {
118 t.Fatalf("Gemini request: %+v", wire)
119 }
120 case <-time.After(3 * time.Second):
121 t.Fatal("extension command did not send Gemini request")
122 }
123}
124
125// This exercises the interactive prelude, core transcript, Gemini request
126// serialization, streaming, and session replay without contacting Google.
127func TestToolPromptGeminiLocalRoundTripAndResume(t *testing.T) {
128 requests := make(chan []byte, 2)
129 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
130 defer r.Body.Close()
131 var raw json.RawMessage
132 if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
133 http.Error(w, err.Error(), http.StatusBadRequest)
134 return
135 }
136 requests <- raw
137 w.Header().Set("Content-Type", "text/event-stream")
138 _, _ = fmt.Fprint(w, "data: "+`{"candidates":[{"content":{"role":"model","parts":[{"text":"done"}]},"finishReason":"STOP"}]}`+"\n\n")
139 }))
140 defer server.Close()
141
142 client := provider.NewGemini("synthetic-test-key", server.URL)
143 agent := core.NewAgent(client, "gemini-3-flash-preview", "", core.Registry{"skill": &shortcutTool{}})
144 session, err := core.NewSession(t.TempDir(), t.TempDir(), "google", "gemini-3-flash-preview", "test")
145 if err != nil {
146 t.Fatal(err)
147 }
148 agent.OnMessageAppended = func(m provider.Message) {
149 if err := session.AppendMessage(m); err != nil {
150 t.Error(err)
151 }
152 }
153 iv := NewInteractive(InteractiveConfig{Agent: agent, Extensions: &extensions.Manager{}})
154 ctx, cancel := context.WithCancel(context.Background())
155 defer cancel()
156 iv.startTurnWithPrelude(ctx, "and push", nil, &toolPromptRequest{
157 call: provider.ToolCallBlock{ID: "ext-prompt-local", Name: "skill", Arguments: json.RawMessage(`{"name":"developer:commit"}`)}, origin: "shortcut",
158 })
159
160 var wire struct {
161 Contents []struct {
162 Role string `json:"role"`
163 Parts []struct {
164 Text string `json:"text"`
165 Signature string `json:"thoughtSignature"`
166 FunctionCall *struct {
167 Name string `json:"name"`
168 } `json:"functionCall"`
169 FunctionResponse *struct {
170 Name string `json:"name"`
171 Response map[string]string `json:"response"`
172 } `json:"functionResponse"`
173 } `json:"parts"`
174 } `json:"contents"`
175 }
176 select {
177 case raw := <-requests:
178 if err := json.Unmarshal(raw, &wire); err != nil {
179 t.Fatal(err)
180 }
181 case <-time.After(3 * time.Second):
182 t.Fatal("Gemini request not sent")
183 }
184 if len(wire.Contents) != 3 || wire.Contents[0].Parts[0].Text != "and push" || wire.Contents[1].Role != "model" || wire.Contents[1].Parts[0].FunctionCall == nil || wire.Contents[1].Parts[0].FunctionCall.Name != "skill" || wire.Contents[1].Parts[0].Signature != "skip_thought_signature_validator" || wire.Contents[2].Parts[0].FunctionResponse == nil || wire.Contents[2].Parts[0].FunctionResponse.Response["output"] != "instructions" {
185 t.Fatalf("Gemini request did not contain one matched synthetic tool exchange: %+v", wire)
186 }
187 deadline := time.NewTimer(3 * time.Second)
188 defer deadline.Stop()
189 poll := time.NewTicker(5 * time.Millisecond)
190 defer poll.Stop()
191 for {
192 iv.mu.Lock()
193 busy := iv.busy
194 iv.mu.Unlock()
195 if !busy {
196 break
197 }
198 select {
199 case <-poll.C:
200 case <-deadline.C:
201 t.Fatal("Gemini turn did not finish")
202 }
203 }
204 path := session.Path
205 if err := session.Close(); err != nil {
206 t.Fatal(err)
207 }
208 reopened, msgs, err := core.OpenSession(path)
209 if err != nil {
210 t.Fatal(err)
211 }
212 t.Cleanup(func() { _ = reopened.Close() })
213 if len(msgs) != 4 || msgs[1].Meta["synthetic_tool_call"] != "true" {
214 t.Fatalf("restored transcript: %+v", msgs)
215 }
216 resumed := core.NewAgent(client, "gemini-3-flash-preview", "", core.Registry{"skill": &shortcutTool{}})
217 resumed.SetMessages(msgs)
218 if err := resumed.Prompt(ctx, "continue", nil, nil); err != nil {
219 t.Fatal(err)
220 }
221 select {
222 case raw := <-requests:
223 if err := json.Unmarshal(raw, &wire); err != nil {
224 t.Fatal(err)
225 }
226 case <-time.After(3 * time.Second):
227 t.Fatal("resumed Gemini request not sent")
228 }
229 if len(wire.Contents) != 5 || wire.Contents[1].Parts[0].Signature != "skip_thought_signature_validator" || wire.Contents[2].Parts[0].FunctionResponse.Response["output"] != "instructions" || wire.Contents[4].Parts[0].Text != "continue" {
230 t.Fatalf("Gemini replay: %+v", wire)
231 }
232}
233
234func TestToolPromptSurvivesPreTurnCompaction(t *testing.T) {
235 client := &compactQueueClient{
236 compactionStarted: make(chan struct{}),
237 releaseCompaction: make(chan struct{}),
238 followUpRequest: make(chan provider.Request, 1),
239 }
240 agent := core.NewAgent(client, "test-model", "", core.Registry{"skill": &shortcutTool{}})
241 agent.SetMessages([]provider.Message{
242 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
243 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
244 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
245 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
246 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
247 })
248 threshold := 70
249 iv := NewInteractive(InteractiveConfig{Agent: agent, Extensions: &extensions.Manager{}, Provider: "anthropic", Model: "claude-sonnet-4-5", AutoCompactThreshold: &threshold})
250 iv.runCtx = context.Background()
251 iv.lastCtxInput = 150000
252 iv.startTurnWithPrelude(context.Background(), "and push", nil, &toolPromptRequest{call: provider.ToolCallBlock{ID: "ext-prompt-2", Name: "skill", Arguments: json.RawMessage(`{}`)}, origin: "shortcut"})
253 select {
254 case <-client.compactionStarted:
255 case <-time.After(2 * time.Second):
256 t.Fatal("compaction did not start")
257 }
258 close(client.releaseCompaction)
259 select {
260 case req := <-client.followUpRequest:
261 if requestUserTextCount(req, "and push") != 1 || len(req.Messages) < 3 || req.Messages[len(req.Messages)-2].Content[0].(provider.ToolCallBlock).Name != "skill" || req.Messages[len(req.Messages)-1].Content[0].(provider.ToolResultBlock).Content[0].(provider.TextBlock).Text != "instructions" {
262 t.Fatalf("prelude lost after compaction: %+v", req.Messages)
263 }
264 case <-time.After(2 * time.Second):
265 t.Fatal("model request not sent")
266 }
267}
268
269func TestToolPromptStartsModelAfterToolResult(t *testing.T) {
270 client := &toolPromptClient{requests: make(chan provider.Request, 1)}
271 agent := core.NewAgent(client, "test-model", "", core.Registry{"skill": &shortcutTool{}})
272 iv := NewInteractive(InteractiveConfig{Agent: agent, Extensions: &extensions.Manager{}})
273 ctx, cancel := context.WithCancel(context.Background())
274 defer cancel()
275 iv.startTurnWithPrelude(ctx, "and push", nil, &toolPromptRequest{
276 call: provider.ToolCallBlock{ID: "ext-prompt-1", Name: "skill", Arguments: json.RawMessage(`{"name":"developer:commit"}`)}, origin: "shortcut",
277 })
278 select {
279 case req := <-client.requests:
280 if len(req.Messages) != 3 || req.Messages[0].Role != provider.RoleUser || req.Messages[1].Role != provider.RoleAssistant || req.Messages[2].Role != provider.RoleTool || req.Messages[2].Content[0].(provider.ToolResultBlock).Content[0].(provider.TextBlock).Text != "instructions" {
281 t.Fatalf("model input: %+v", req.Messages)
282 }
283 case <-time.After(2 * time.Second):
284 t.Fatal("model request not sent")
285 }
286}