1package core
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "strings"
8 "sync"
9 "sync/atomic"
10 "testing"
11 "time"
12
13 "github.com/bnema/zut/packages/provider"
14)
15
16// continuationScriptClient replays a scripted stop reason per provider call
17// and records every request for identity/history assertions.
18type continuationScriptClient struct {
19 mu sync.Mutex
20 requests []provider.Request
21 calls int32
22 // continuations is the number of leading calls that return
23 // StopContinue before one final StopEnd.
24 continuations int32
25}
26
27func (c *continuationScriptClient) Name() string { return "continuation-script" }
28
29func (c *continuationScriptClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
30 call := atomic.AddInt32(&c.calls, 1)
31 c.mu.Lock()
32 c.requests = append(c.requests, req)
33 c.mu.Unlock()
34 out := make(chan provider.Event, 4)
35 out <- provider.EventStart{Provider: c.Name(), Model: req.Model}
36 text := "final"
37 stop := provider.StopEnd
38 if call <= c.continuations {
39 text = fmt.Sprintf("part-%d", call)
40 stop = provider.StopContinue
41 }
42 out <- provider.EventTextDelta{Delta: text}
43 out <- provider.EventUsage{Usage: provider.Usage{InputTokens: 1, OutputTokens: 1}}
44 out <- provider.EventDone{Stop: stop, Message: provider.Message{
45 Role: provider.RoleAssistant,
46 Content: []provider.Content{provider.TextBlock{Text: text}},
47 }}
48 close(out)
49 return out, nil
50}
51
52func (c *continuationScriptClient) Calls() int32 { return atomic.LoadInt32(&c.calls) }
53
54func (c *continuationScriptClient) Requests() []provider.Request {
55 c.mu.Lock()
56 defer c.mu.Unlock()
57 out := make([]provider.Request, len(c.requests))
58 copy(out, c.requests)
59 return out
60}
61
62type continuationEvents struct {
63 mu sync.Mutex
64 turnEnds []EvTurnEnd
65 doneCount int
66}
67
68func (e *continuationEvents) sink(ev AgentEvent) {
69 e.mu.Lock()
70 defer e.mu.Unlock()
71 if te, ok := ev.(EvTurnEnd); ok {
72 e.turnEnds = append(e.turnEnds, te)
73 }
74 if _, ok := ev.(EvDone); ok {
75 e.doneCount++
76 }
77}
78
79func (e *continuationEvents) snapshot() ([]EvTurnEnd, int) {
80 e.mu.Lock()
81 defer e.mu.Unlock()
82 out := make([]EvTurnEnd, len(e.turnEnds))
83 copy(out, e.turnEnds)
84 return out, e.doneCount
85}
86
87// TestAgentContinuationTwelveStepsThenDone pins AC3: 12 text-only explicit
88// continuations followed by a final response produce 13 requests and one
89// terminal completion, with no artificial user/tool rows.
90func TestAgentContinuationTwelveStepsThenDone(t *testing.T) {
91 client := &continuationScriptClient{continuations: 12}
92 agent := NewAgent(client, "test-model", "system", Registry{})
93 events := &continuationEvents{}
94
95 ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
96 defer cancel()
97 if err := agent.Prompt(ctx, "start", nil, events.sink); err != nil {
98 t.Fatalf("Prompt returned %v", err)
99 }
100 if got := client.Calls(); got != 13 {
101 t.Fatalf("provider calls = %d, want 13", got)
102 }
103 turnEnds, doneCount := events.snapshot()
104 if doneCount != 1 {
105 t.Fatalf("terminal EvDone count = %d, want 1", doneCount)
106 }
107 if len(turnEnds) != 13 {
108 t.Fatalf("EvTurnEnd count = %d, want 13", len(turnEnds))
109 }
110 for i, te := range turnEnds {
111 want := provider.StopContinue
112 if i == 12 {
113 want = provider.StopEnd
114 }
115 if te.Stop != want || te.Err != nil {
116 t.Fatalf("turn %d end = (%q, %v), want (%q, nil)", i+1, te.Stop, te.Err, want)
117 }
118 }
119
120 msgs := agent.Messages()
121 if len(msgs) != 14 {
122 t.Fatalf("transcript messages = %d, want 1 user + 13 assistant", len(msgs))
123 }
124 if msgs[0].Role != provider.RoleUser {
125 t.Fatalf("message 0 role = %q, want user", msgs[0].Role)
126 }
127 for i := 1; i < len(msgs); i++ {
128 if msgs[i].Role != provider.RoleAssistant {
129 t.Fatalf("message %d role = %q, want assistant (no artificial rows)", i, msgs[i].Role)
130 }
131 if len(msgs[i].Content) == 0 {
132 t.Fatalf("message %d has no content", i)
133 }
134 }
135
136 // Turn/thread identity is stable across the whole logical turn and
137 // every continuation sees prior assistant content.
138 requests := client.Requests()
139 turnID := requests[0].Context.TurnID
140 if turnID == "" {
141 t.Fatal("first request has empty turn ID")
142 }
143 for i, req := range requests {
144 if req.Context.TurnID != turnID {
145 t.Fatalf("request %d turn ID = %q, want stable %q", i, req.Context.TurnID, turnID)
146 }
147 if want := 1 + i; len(req.Messages) != want {
148 t.Fatalf("request %d messages = %d, want %d (prior content retained)", i, len(req.Messages), want)
149 }
150 }
151
152 // Usage is counted once per inference.
153 if got := agent.Cost(); got.InputTokens != 13 || got.OutputTokens != 13 {
154 t.Fatalf("cumulative usage = %+v, want 13 input / 13 output", got)
155 }
156 if got := extractText(msgs[len(msgs)-1]); got != "final" {
157 t.Fatalf("final text = %q, want %q", got, "final")
158 }
159}
160
161// TestAgentContinuationCanceledAtBoundary pins the canceled path: no
162// further inference is requested and the cancellation error surfaces.
163func TestAgentContinuationCanceledAtBoundary(t *testing.T) {
164 client := &continuationScriptClient{continuations: 12}
165 agent := NewAgent(client, "test-model", "system", Registry{})
166 events := &continuationEvents{}
167
168 ctx, cancel := context.WithCancel(context.Background())
169 defer cancel()
170 var once sync.Once
171 sink := func(ev AgentEvent) {
172 events.sink(ev)
173 if te, ok := ev.(EvTurnEnd); ok && te.Stop == provider.StopContinue {
174 once.Do(cancel)
175 }
176 }
177 err := agent.Prompt(ctx, "start", nil, sink)
178 if !errors.Is(err, context.Canceled) {
179 t.Fatalf("Prompt error = %v, want context.Canceled", err)
180 }
181 if got := client.Calls(); got != 1 {
182 t.Fatalf("provider calls after cancel = %d, want 1", got)
183 }
184}
185
186// TestAgentContinuationRespectsMaxSteps pins that configured step limits
187// still bound continuations with the existing typed error.
188func TestAgentContinuationRespectsMaxSteps(t *testing.T) {
189 client := &continuationScriptClient{continuations: 12}
190 agent := NewAgent(client, "test-model", "system", Registry{})
191 agent.MaxSteps = 2
192
193 err := agent.Prompt(context.Background(), "start", nil, nil)
194 if !errors.Is(err, ErrMaxSteps) {
195 t.Fatalf("Prompt error = %v, want ErrMaxSteps", err)
196 }
197 if got := client.Calls(); got != 2 {
198 t.Fatalf("provider calls = %d, want 2", got)
199 }
200}
201
202// TestAgentContinuationGuardBlocksNextRequest pins that turn guards stay
203// active across the continuation boundary.
204func TestAgentContinuationGuardBlocksNextRequest(t *testing.T) {
205 client := &continuationScriptClient{continuations: 12}
206 agent := NewAgent(client, "test-model", "system", Registry{})
207 agent.BeforeTurn = func(step int) (bool, string) {
208 if step > 1 {
209 return false, "blocked at step 2"
210 }
211 return true, ""
212 }
213 events := &continuationEvents{}
214
215 if err := agent.Prompt(context.Background(), "start", nil, events.sink); err != nil {
216 t.Fatalf("Prompt returned %v", err)
217 }
218 if got := client.Calls(); got != 1 {
219 t.Fatalf("provider calls = %d, want 1 (guard blocks the next request)", got)
220 }
221 turnEnds, doneCount := events.snapshot()
222 if doneCount != 1 {
223 t.Fatalf("terminal EvDone count = %d, want 1", doneCount)
224 }
225 if len(turnEnds) != 2 {
226 t.Fatalf("EvTurnEnd count = %d, want continuation + guard", len(turnEnds))
227 }
228 if turnEnds[0].Stop != provider.StopContinue {
229 t.Fatalf("first turn end = %q, want %q", turnEnds[0].Stop, provider.StopContinue)
230 }
231 if turnEnds[1].Stop != provider.StopError || turnEnds[1].Err == nil {
232 t.Fatalf("guard turn end = (%q, %v), want (error, non-nil)", turnEnds[1].Stop, turnEnds[1].Err)
233 }
234}
235
236// TestAgentContinuationConsumesQueuedInput pins that user text queued
237// during a continuation is delivered at the normal boundary, answered by
238// the next inference, and never synthesized by the loop itself.
239func TestAgentContinuationConsumesQueuedInput(t *testing.T) {
240 client := &continuationScriptClient{continuations: 1}
241 agent := NewAgent(client, "test-model", "system", Registry{})
242 events := &continuationEvents{}
243
244 var once sync.Once
245 sink := func(ev AgentEvent) {
246 events.sink(ev)
247 if te, ok := ev.(EvTurnEnd); ok && te.Stop == provider.StopContinue {
248 once.Do(func() {
249 if !agent.QueueMessage("follow-up", nil) {
250 t.Error("QueueMessage returned false")
251 }
252 })
253 }
254 }
255 ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
256 defer cancel()
257 if err := agent.Prompt(ctx, "start", nil, sink); err != nil {
258 t.Fatalf("Prompt returned %v", err)
259 }
260 if got := client.Calls(); got != 2 {
261 t.Fatalf("provider calls = %d, want 2", got)
262 }
263 requests := client.Requests()
264 if len(requests) != 2 {
265 t.Fatalf("recorded requests = %d, want 2", len(requests))
266 }
267 var sawQueued bool
268 for _, m := range requests[1].Messages {
269 if m.Role == provider.RoleUser && strings.Contains(extractText(m), "follow-up") {
270 sawQueued = true
271 }
272 }
273 if !sawQueued {
274 t.Fatalf("second request did not carry queued input: %#v", requests[1].Messages)
275 }
276 msgs := agent.Messages()
277 var userCount int
278 for _, m := range msgs {
279 if m.Role == provider.RoleUser {
280 userCount++
281 }
282 }
283 if userCount != 2 {
284 t.Fatalf("user messages = %d, want initial + queued", userCount)
285 }
286 if _, doneCount := events.snapshot(); doneCount != 1 {
287 t.Fatalf("terminal EvDone count = %d, want 1", doneCount)
288 }
289}