1package modes
2
3import (
4 "context"
5 "strings"
6 "testing"
7 "time"
8
9 "github.com/patriceckhart/zot/packages/core"
10 "github.com/patriceckhart/zot/packages/provider"
11)
12
13type preparationCaptureClient struct{ requests chan provider.Request }
14
15func (*preparationCaptureClient) Name() string { return "test" }
16func (c *preparationCaptureClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
17 select {
18 case c.requests <- req:
19 case <-ctx.Done():
20 return nil, ctx.Err()
21 }
22 events := make(chan provider.Event, 1)
23 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "ok"}}}}
24 close(events)
25 return events, nil
26}
27
28func TestAutoSwarmToggleDuringPromptPreparation(t *testing.T) {
29 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
30 defer cancel()
31 client := &preparationCaptureClient{requests: make(chan provider.Request, 1)}
32 ag := core.NewAgent(client, "model", "base", nil)
33 entered, release, done := make(chan struct{}), make(chan struct{}), make(chan struct{})
34 calls := 0
35 ag.BeforeStart = func(ctx context.Context, system string) string {
36 calls++
37 if calls == 1 {
38 close(entered)
39 select {
40 case <-release:
41 case <-ctx.Done():
42 }
43 return "stale replacement"
44 }
45 return system + "/extension"
46 }
47 i := &Interactive{agent: ag, cfg: InteractiveConfig{AutoSwarmSystemAddendum: "swarm instructions"}}
48 var runErr error
49 go func() { defer close(done); runErr = ag.Prompt(ctx, "hello", nil, nil) }()
50 t.Cleanup(func() { cancel(); <-done })
51 select {
52 case <-entered:
53 case <-ctx.Done():
54 t.Fatal("hook did not start")
55 }
56 i.applyAutoSwarmSystemPrompt(true)
57 close(release)
58 select {
59 case <-done:
60 case <-ctx.Done():
61 t.Fatal("prompt did not finish")
62 }
63 if runErr != nil {
64 t.Fatal(runErr)
65 }
66 if calls != 2 {
67 t.Fatalf("preparation calls = %d, want 2", calls)
68 }
69 select {
70 case req := <-client.requests:
71 if req.System != "base\n\nswarm instructions/extension" {
72 t.Fatalf("provider received stale prompt: %q", req.System)
73 }
74 default:
75 t.Fatal("provider was not called")
76 }
77 if ag.BaseSystem() != "base\n\nswarm instructions" {
78 t.Fatal("rebuild base was lost")
79 }
80}
81
82func TestAutoSwarmToggleAfterPreparationBeforeModelCall(t *testing.T) {
83 client := &preparationCaptureClient{requests: make(chan provider.Request, 1)}
84 ag := core.NewAgent(client, "model", "base", nil)
85 i := &Interactive{agent: ag, cfg: InteractiveConfig{AutoSwarmSystemAddendum: "swarm instructions"}}
86 calls := 0
87 ag.BeforeStart = func(_ context.Context, system string) string {
88 calls++
89 return system + "/extension"
90 }
91 // Simulate a settings change while the turn-start guard runs, after
92 // initial preparation but before the provider request is constructed.
93 ag.BeforeTurn = func(int) (bool, string) {
94 i.applyAutoSwarmSystemPrompt(true)
95 return true, ""
96 }
97 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
98 defer cancel()
99 if err := ag.Prompt(ctx, "hello", nil, nil); err != nil {
100 t.Fatal(err)
101 }
102 if calls != 2 {
103 t.Fatalf("preparation calls = %d, want 2", calls)
104 }
105 select {
106 case req := <-client.requests:
107 if req.System != "base\n\nswarm instructions/extension" {
108 t.Fatalf("provider received unprepared prompt: %q", req.System)
109 }
110 default:
111 t.Fatal("provider was not called")
112 }
113}
114
115func TestAutoSwarmRebuildsBasePromptBeforeExtensions(t *testing.T) {
116 ag := core.NewAgent(nil, "model", "base", nil)
117 calls := 0
118 ag.BeforeStart = func(ctx context.Context, system string) string {
119 calls++
120 if strings.Contains(system, "extension") {
121 t.Fatal("replacement fed back into extension")
122 }
123 return system + "/extension"
124 }
125 ag.BeforeTurn = func(int) (bool, string) { return false, "test" }
126 i := &Interactive{agent: ag, cfg: InteractiveConfig{AutoSwarmSystemAddendum: "swarm instructions"}}
127 prepare := func() {
128 t.Helper()
129 if err := ag.Continue(context.Background(), nil); err != nil {
130 t.Fatal(err)
131 }
132 }
133 prepare()
134 i.applyAutoSwarmSystemPrompt(true)
135 prepare()
136 if !strings.Contains(ag.System, "swarm instructions") || calls != 2 {
137 t.Fatal("toggle did not rebuild prompt")
138 }
139 i.applyAutoSwarmSystemPrompt(true)
140 prepare()
141 if calls != 2 {
142 t.Fatal("unchanged toggle reapplied extensions")
143 }
144 i.applyAutoSwarmSystemPrompt(false)
145 prepare()
146 if strings.Contains(ag.System, "swarm instructions") || calls != 3 {
147 t.Fatal("toggle did not remove instructions")
148 }
149}