1package core
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "strings"
9 "sync/atomic"
10 "testing"
11
12 "github.com/bnema/zut/packages/provider"
13)
14
15type repetitionTestClient struct {
16 calls atomic.Int32
17 assistantOnly bool
18}
19
20func (c *repetitionTestClient) Name() string { return "repetition-test" }
21
22func (c *repetitionTestClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
23 call := c.calls.Add(1)
24 id := fmt.Sprintf("call-%d", call)
25 out := make(chan provider.Event, 4)
26 out <- provider.EventStart{Provider: c.Name(), Model: req.Model}
27 if c.assistantOnly {
28 out <- provider.EventTextDelta{Delta: "I am still checking."}
29 out <- provider.EventDone{Stop: provider.StopContinue, Message: provider.Message{
30 Role: provider.RoleAssistant,
31 Content: []provider.Content{provider.TextBlock{Text: "I am still checking."}},
32 }}
33 } else {
34 out <- provider.EventToolStart{ID: id, Name: "repeat"}
35 out <- provider.EventToolEnd{ID: id}
36 out <- provider.EventDone{Stop: provider.StopToolUse, Message: provider.Message{
37 Role: provider.RoleAssistant,
38 Content: []provider.Content{provider.ToolCallBlock{
39 ID: id, Name: "repeat", Arguments: json.RawMessage(`{"b":2,"a":1}`),
40 }},
41 }}
42 }
43 close(out)
44 return out, nil
45}
46
47type repetitionTestTool struct {
48 calls atomic.Int32
49 differ bool
50}
51
52func (*repetitionTestTool) Name() string { return "repeat" }
53func (*repetitionTestTool) Description() string { return "repeats a test operation" }
54func (*repetitionTestTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
55func (t *repetitionTestTool) Execute(context.Context, json.RawMessage, func(string)) (ToolResult, error) {
56 call := t.calls.Add(1)
57 result := "unchanged"
58 if t.differ {
59 result = fmt.Sprintf("result %d", call)
60 }
61 return ToolResult{Content: []provider.Content{provider.TextBlock{Text: result}}}, nil
62}
63
64func TestAgentStopsAfterWarningForRepeatedToolCallAndResult(t *testing.T) {
65 client := &repetitionTestClient{}
66 tool := &repetitionTestTool{}
67 agent := NewAgent(client, "model", "system", Registry{"repeat": tool})
68 var warnings, stops int
69 if err := agent.Prompt(context.Background(), "inspect the project", nil, func(event AgentEvent) {
70 guard, ok := event.(EvRepetitionGuard)
71 if !ok {
72 return
73 }
74 switch guard.Stage {
75 case RepetitionGuardWarning:
76 warnings++
77 if guard.Count != repetitionWarnThreshold {
78 t.Errorf("warning count = %d, want %d", guard.Count, repetitionWarnThreshold)
79 }
80 case RepetitionGuardStopped:
81 stops++
82 if guard.Count != repetitionStopThreshold {
83 t.Errorf("stop count = %d, want %d", guard.Count, repetitionStopThreshold)
84 }
85 }
86 }); !errors.Is(err, ErrRepetitiveLoop) {
87 t.Fatalf("Prompt error = %v, want ErrRepetitiveLoop", err)
88 }
89 if got := client.calls.Load(); got != int32(repetitionStopThreshold) {
90 t.Fatalf("provider calls = %d, want %d", got, repetitionStopThreshold)
91 }
92 if got := warnings; got != 1 {
93 t.Fatalf("warning events = %d, want 1", got)
94 }
95 if got := stops; got != 1 {
96 t.Fatalf("stop events = %d, want 1", got)
97 }
98
99 messages := agent.Messages()
100 var warningFound bool
101 for _, message := range messages {
102 if message.Meta[RepetitionGuardMetaKey] == "true" {
103 warningFound = true
104 if message.Role != provider.RoleUser || extractText(message) == "" {
105 t.Fatalf("malformed hidden warning message: %#v", message)
106 }
107 }
108 }
109 if !warningFound {
110 t.Fatal("warning was not added to the transcript")
111 }
112 if len(messages) < 3 {
113 t.Fatalf("transcript has %d messages, want tool-call history", len(messages))
114 }
115 if got := messages[len(messages)-1].Role; got != provider.RoleTool {
116 t.Fatalf("last transcript role = %q, want paired tool result", got)
117 }
118 if got := messages[len(messages)-2].Content[0].(provider.ToolCallBlock).Name; got != "repeat" {
119 t.Fatalf("repeated tool name = %q, want repeat", got)
120 }
121 if !strings.Contains(repetitionGuardPromptFromTranscript(messages), "tool \"repeat\"") {
122 t.Fatal("warning prompt does not identify the repeated tool")
123 }
124}
125
126func repetitionGuardPromptFromTranscript(messages []provider.Message) string {
127 for _, message := range messages {
128 if message.Meta[RepetitionGuardMetaKey] == "true" {
129 return extractText(message)
130 }
131 }
132 return ""
133}
134
135func TestToolCallResultFingerprintCanonicalizesJSONArgumentKeyOrder(t *testing.T) {
136 first := provider.ToolCallBlock{Name: "repeat", Arguments: json.RawMessage(`{"b":2,"a":1}`)}
137 second := provider.ToolCallBlock{Name: "repeat", Arguments: json.RawMessage(`{"a":1,"b":2}`)}
138 result := provider.ToolResultBlock{Content: []provider.Content{provider.TextBlock{Text: "same"}}}
139 if firstFingerprint, secondFingerprint := toolCallResultFingerprint(first, result), toolCallResultFingerprint(second, result); firstFingerprint != secondFingerprint {
140 t.Fatalf("fingerprints differ for equivalent JSON args: %q != %q", firstFingerprint, secondFingerprint)
141 }
142}
143
144func TestAgentDoesNotStopWhenRepeatedToolOutputChanges(t *testing.T) {
145 client := &repetitionTestClient{}
146 tool := &repetitionTestTool{differ: true}
147 agent := NewAgent(client, "model", "system", Registry{"repeat": tool})
148 agent.MaxSteps = repetitionStopThreshold + 1
149 var guardEvents int
150 err := agent.Prompt(context.Background(), "inspect the project", nil, func(event AgentEvent) {
151 if _, ok := event.(EvRepetitionGuard); ok {
152 guardEvents++
153 }
154 })
155 if errors.Is(err, ErrRepetitiveLoop) {
156 t.Fatalf("Prompt error = %v, unexpected loop detection for changing results", err)
157 }
158 if guardEvents != 0 {
159 t.Fatalf("repetition guard events = %d, want 0", guardEvents)
160 }
161}
162
163func TestAgentCompactionRollbackPreservesRepetitionCount(t *testing.T) {
164 client := &repetitionTestClient{}
165 agent := NewAgent(client, "model", "system", Registry{})
166 for i := 0; i < repetitionWarnThreshold-1; i++ {
167 agent.observeRepetition("repeated", RepetitionKindAssistantMessage, "")
168 }
169 before := agent.repetition.patterns["repeated"].count
170 if before != repetitionWarnThreshold-1 {
171 t.Fatalf("repetition count before rollback = %d, want %d", before, repetitionWarnThreshold-1)
172 }
173
174 agent.SetMessages([]provider.Message{{Role: provider.RoleUser}})
175 agent.RestoreMessages(nil)
176 if got := agent.repetition.patterns["repeated"].count; got != before {
177 t.Fatalf("repetition count after transcript rollback = %d, want %d", got, before)
178 }
179}
180
181func TestAgentTranscriptReplacementResetsRepetitionGuard(t *testing.T) {
182 client := &repetitionTestClient{}
183 agent := NewAgent(client, "model", "system", Registry{})
184 agent.DisableRepetitionGuard = true
185 for i := 0; i < repetitionStopThreshold; i++ {
186 agent.observeRepetition("stopped", RepetitionKindAssistantMessage, "")
187 }
188 agent.repetition.stopped = true
189 agent.repetition.stoppedDecision = repetitionDecision{count: repetitionStopThreshold}
190
191 agent.SetMessages(nil)
192 agent.ResetRepetitionGuard()
193 agent.DisableRepetitionGuard = false
194 agent.MaxSteps = 1
195 if err := agent.Continue(context.Background(), nil); !errors.Is(err, ErrMaxSteps) {
196 t.Fatalf("Continue after transcript replacement error = %v, want ErrMaxSteps after one request", err)
197 }
198 if got := client.calls.Load(); got != 1 {
199 t.Fatalf("provider calls = %d, want 1", got)
200 }
201}
202
203func TestAgentResetRepetitionGuardAllowsContinueAfterStop(t *testing.T) {
204 client := &repetitionTestClient{}
205 agent := NewAgent(client, "model", "system", Registry{})
206 agent.DisableRepetitionGuard = true
207 agent.MaxSteps = 1
208 for i := 0; i < repetitionStopThreshold; i++ {
209 agent.observeRepetition("stopped", RepetitionKindAssistantMessage, "")
210 }
211 agent.repetition.stopped = true
212 agent.repetition.stoppedDecision = repetitionDecision{toolName: "repeat", count: repetitionStopThreshold}
213 if agent.repetitionStopped() {
214 t.Fatal("repetition guard should be disabled")
215 }
216
217 agent.ResetRepetitionGuard()
218 agent.DisableRepetitionGuard = false
219 if err := agent.Continue(context.Background(), nil); !errors.Is(err, ErrMaxSteps) {
220 t.Fatalf("Continue after reset error = %v, want ErrMaxSteps after one request", err)
221 }
222 if got := client.calls.Load(); got != 1 {
223 t.Fatalf("provider calls = %d, want 1", got)
224 }
225}
226
227func TestAgentRepetitionCountSurvivesContinue(t *testing.T) {
228 client := &repetitionTestClient{}
229 tool := &repetitionTestTool{}
230 agent := NewAgent(client, "model", "system", Registry{"repeat": tool})
231 agent.MaxSteps = 4
232 if err := agent.Prompt(context.Background(), "inspect the project", nil, nil); !errors.Is(err, ErrMaxSteps) {
233 t.Fatalf("Prompt error = %v, want ErrMaxSteps after four calls", err)
234 }
235
236 var warningCount, stoppedCount int
237 err := agent.Continue(context.Background(), func(event AgentEvent) {
238 guard, ok := event.(EvRepetitionGuard)
239 if !ok {
240 return
241 }
242 if guard.Stage == RepetitionGuardWarning {
243 warningCount = guard.Count
244 }
245 if guard.Stage == RepetitionGuardStopped {
246 stoppedCount = guard.Count
247 }
248 })
249 if !errors.Is(err, ErrRepetitiveLoop) {
250 t.Fatalf("Continue error = %v, want ErrRepetitiveLoop", err)
251 }
252 if warningCount != repetitionWarnThreshold || stoppedCount != repetitionStopThreshold {
253 t.Fatalf("warning/stop counts = %d/%d, want %d/%d", warningCount, stoppedCount, repetitionWarnThreshold, repetitionStopThreshold)
254 }
255 if got := client.calls.Load(); got != int32(repetitionStopThreshold) {
256 t.Fatalf("provider calls = %d, want %d", got, repetitionStopThreshold)
257 }
258}
259
260func TestRepetitionGuardWindowEvictsOldPatterns(t *testing.T) {
261 agent := NewAgent(nil, "model", "system", nil)
262 fingerprint := strings.Repeat("a", 64)
263 for range repetitionWarnThreshold {
264 decision := agent.observeRepetition(fingerprint, RepetitionKindToolCall, "read")
265 if decision.stage == RepetitionGuardWarning {
266 break
267 }
268 }
269 for index := 0; index < repetitionWindowSize; index++ {
270 agent.observeRepetition(fmt.Sprintf("other-%d", index), RepetitionKindAssistantMessage, "")
271 }
272 for index := 0; index < repetitionWarnThreshold-1; index++ {
273 if decision := agent.observeRepetition(fingerprint, RepetitionKindToolCall, "read"); decision.stage != "" {
274 t.Fatalf("repetition %d after eviction reported stage %q", index+1, decision.stage)
275 }
276 }
277 decision := agent.observeRepetition(fingerprint, RepetitionKindToolCall, "read")
278 if decision.stage != RepetitionGuardWarning || decision.count != repetitionWarnThreshold {
279 t.Fatalf("post-eviction decision = %#v, want warning at %d", decision, repetitionWarnThreshold)
280 }
281}
282
283func TestAgentDisableRepetitionGuardAllowsPolling(t *testing.T) {
284 client := &repetitionTestClient{}
285 tool := &repetitionTestTool{}
286 agent := NewAgent(client, "model", "system", Registry{"repeat": tool})
287 agent.DisableRepetitionGuard = true
288 agent.MaxSteps = repetitionStopThreshold + 1
289 if err := agent.Prompt(context.Background(), "poll", nil, nil); !errors.Is(err, ErrMaxSteps) {
290 t.Fatalf("Prompt error = %v, want MaxSteps with the guard disabled", err)
291 }
292 if got := client.calls.Load(); got != repetitionStopThreshold+1 {
293 t.Fatalf("provider calls = %d, want %d", got, repetitionStopThreshold+1)
294 }
295}
296
297func TestAgentStopsAfterWarningForRepeatedAssistantMessage(t *testing.T) {
298 client := &repetitionTestClient{assistantOnly: true}
299 agent := NewAgent(client, "model", "system", Registry{})
300 var warnings, stops int
301 err := agent.Prompt(context.Background(), "continue the goal", nil, func(event AgentEvent) {
302 guard, ok := event.(EvRepetitionGuard)
303 if !ok {
304 return
305 }
306 if guard.Kind != RepetitionKindAssistantMessage {
307 t.Errorf("guard kind = %q, want repeated assistant message", guard.Kind)
308 }
309 if guard.Stage == RepetitionGuardWarning {
310 warnings++
311 }
312 if guard.Stage == RepetitionGuardStopped {
313 stops++
314 }
315 })
316 if !errors.Is(err, ErrRepetitiveLoop) {
317 t.Fatalf("Prompt error = %v, want ErrRepetitiveLoop", err)
318 }
319 if got := client.calls.Load(); got != int32(repetitionStopThreshold) {
320 t.Fatalf("provider calls = %d, want %d", got, repetitionStopThreshold)
321 }
322 if warnings != 1 || stops != 1 {
323 t.Fatalf("warning/stop events = %d/%d, want 1/1", warnings, stops)
324 }
325}