1package modes
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "errors"
8 "strings"
9 "sync"
10 "testing"
11 "time"
12
13 "github.com/bnema/zut/packages/core"
14 "github.com/bnema/zut/packages/provider"
15)
16
17// turnRecoveryModeClient replays one assistant message per Stream call so
18// mode tests can drive the core recovery allowance without timing sleeps.
19type turnRecoveryModeClient struct {
20 mu sync.Mutex
21 calls int
22 turns [][]provider.Content
23}
24
25func (c *turnRecoveryModeClient) Name() string { return "turn-recovery-mode-test" }
26
27func (c *turnRecoveryModeClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
28 c.mu.Lock()
29 call := c.calls
30 c.calls++
31 c.mu.Unlock()
32 var content []provider.Content
33 if call < len(c.turns) {
34 content = c.turns[call]
35 }
36 out := make(chan provider.Event, 4)
37 out <- provider.EventStart{Provider: c.Name(), Model: "test-model"}
38 out <- provider.EventUsage{Usage: provider.Usage{InputTokens: 3, OutputTokens: 5}}
39 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
40 Role: provider.RoleAssistant,
41 Content: content,
42 Time: time.Now(),
43 }}
44 close(out)
45 return out, nil
46}
47
48func (c *turnRecoveryModeClient) Calls() int {
49 c.mu.Lock()
50 defer c.mu.Unlock()
51 return c.calls
52}
53
54func commentaryContent(text string) []provider.Content {
55 return []provider.Content{provider.TextBlock{Text: text, Phase: provider.TextPhaseCommentary}}
56}
57
58func answerContent(text string) []provider.Content {
59 return []provider.Content{provider.TextBlock{Text: text}}
60}
61
62func TestTurnRecoveryModeRendersOneRecoveryStatus(t *testing.T) {
63 client := &turnRecoveryModeClient{turns: [][]provider.Content{
64 commentaryContent("still working on it"),
65 answerContent("finished"),
66 }}
67 ag := core.NewAgent(client, "test-model", "", nil)
68 i := newNotesTestInteractive()
69 i.cfg.Agent = ag
70 i.agent = ag
71
72 i.handleEvent(core.EvTurnRecovery{Reason: core.TurnRecoveryReasonCommentaryOnly, Attempt: 1})
73
74 if len(i.extNotes) != 1 {
75 t.Fatalf("extNotes = %d, want exactly one recovery status", len(i.extNotes))
76 }
77 if got := stripANSIBytes(i.extNotes[0]); !strings.Contains(got, "Continuing: no final answer received.") {
78 t.Fatalf("recovery status = %q, want exact recovery line", got)
79 }
80}
81
82func TestTurnRecoveryPrintSurfacesRecoveryAndCompletes(t *testing.T) {
83 client := &turnRecoveryModeClient{turns: [][]provider.Content{
84 commentaryContent("still working on it"),
85 answerContent("finished"),
86 }}
87 ag := core.NewAgent(client, "test-model", "", nil)
88 var out bytes.Buffer
89 usage, err := RunPrint(context.Background(), ag, "work", nil, &out)
90 if err != nil {
91 t.Fatalf("RunPrint error = %v", err)
92 }
93 _ = usage
94 if out.String() != "finished\n" {
95 t.Fatalf("print output = %q, want recovered final answer", out.String())
96 }
97 if client.Calls() != 2 {
98 t.Fatalf("provider calls = %d, want 2 (one bounded recovery)", client.Calls())
99 }
100}
101
102func TestTurnRecoveryStreamKeepsTranscriptOnStdout(t *testing.T) {
103 client := &turnRecoveryModeClient{turns: [][]provider.Content{
104 nil,
105 answerContent("finished"),
106 }}
107 ag := core.NewAgent(client, "test-model", "", nil)
108 var out, diag bytes.Buffer
109 if err := RunStreamWithDiag(context.Background(), ag, "work", nil, &out, &diag); err != nil {
110 t.Fatalf("RunStreamWithDiag error = %v", err)
111 }
112 if !strings.Contains(out.String(), "finished") {
113 t.Fatalf("stdout = %q, want recovered answer text", out.String())
114 }
115 if !strings.Contains(diag.String(), "Continuing: no final answer received.") {
116 t.Fatalf("diag = %q, want one recovery status", diag.String())
117 }
118 if strings.Contains(out.String(), "Continuing: no final answer received.") {
119 t.Fatalf("stdout = %q, recovery status must not pollute the transcript", out.String())
120 }
121}
122
123func TestTurnRecoveryExhaustionSurfacesIncompleteError(t *testing.T) {
124 client := &turnRecoveryModeClient{turns: [][]provider.Content{
125 nil,
126 nil,
127 }}
128 ag := core.NewAgent(client, "test-model", "", nil)
129 var out bytes.Buffer
130 _, err := RunPrint(context.Background(), ag, "work", nil, &out)
131 if !errors.Is(err, core.ErrIncompleteTurn) {
132 t.Fatalf("RunPrint error = %v, want ErrIncompleteTurn via errors.Is", err)
133 }
134 if out.String() != "" {
135 t.Fatalf("print output = %q, exhaustion must not report success text", out.String())
136 }
137 if client.Calls() != 2 {
138 t.Fatalf("provider calls = %d, want exactly 2 (no unbounded retry)", client.Calls())
139 }
140}
141
142func TestTurnRecoveryGoalDoesNotScheduleDuplicateWork(t *testing.T) {
143 var mu sync.Mutex
144 goal := &core.SessionGoal{ID: "goal-1", Status: core.GoalActive, Objective: "finish the work"}
145 current := func() *core.SessionGoal {
146 mu.Lock()
147 defer mu.Unlock()
148 cp := *goal
149 return &cp
150 }
151 persist := func(next *core.SessionGoal) error {
152 mu.Lock()
153 goal = next
154 mu.Unlock()
155 return nil
156 }
157 client := &turnRecoveryModeClient{turns: [][]provider.Content{
158 commentaryContent("still working on it"),
159 answerContent("finished"),
160 }}
161 ag := core.NewAgent(client, "test-model", "", goalToolRegistry())
162 i := NewInteractive(InteractiveConfig{
163 Agent: ag,
164 CurrentGoal: current,
165 PersistGoal: persist,
166 PersistGoalRuntime: persist,
167 })
168 i.runCtx = context.Background()
169 run, err := i.startGoalRun(current())
170 if err != nil {
171 t.Fatalf("startGoalRun error = %v", err)
172 }
173 if run == nil {
174 t.Fatal("startGoalRun did not start a run")
175 }
176 var events []core.AgentEvent
177 if err := ag.Prompt(context.Background(), "work", nil, func(ev core.AgentEvent) {
178 events = append(events, ev)
179 i.observeGoalRun(ev)
180 i.handleEvent(ev)
181 }); err != nil {
182 t.Fatalf("Prompt error = %v", err)
183 }
184 recoveries := 0
185 for _, ev := range events {
186 if _, ok := ev.(core.EvTurnRecovery); ok {
187 recoveries++
188 }
189 }
190 if recoveries != 1 {
191 t.Fatalf("recovery events = %d, want exactly one inner continuation", recoveries)
192 }
193 // The inner continuation must not start a second goal run: the run
194 // started above is still the single owner, and finishing it once must
195 // record accounting without scheduling another continuation here.
196 i.mu.Lock()
197 owner := i.goalRun
198 i.mu.Unlock()
199 if owner != run {
200 t.Fatal("inner continuation replaced the goal run: want the single owner started above")
201 }
202 if owner.goalID != "goal-1" {
203 t.Fatalf("goal run goalID = %q, want %q", owner.goalID, "goal-1")
204 }
205 if !i.finishGoalRun(false) {
206 t.Fatalf("finishGoalRun = false, want accounting for the single owner run")
207 }
208 mu.Lock()
209 used := goal.TokensUsed
210 mu.Unlock()
211 if used == 0 {
212 t.Fatalf("goal tokens used = 0, want usage recorded once for the single run")
213 }
214}
215
216func TestTurnRecoveryExhaustionKeepsGoalBlocked(t *testing.T) {
217 var mu sync.Mutex
218 goal := &core.SessionGoal{ID: "goal-1", Status: core.GoalActive, Objective: "finish the work"}
219 current := func() *core.SessionGoal {
220 mu.Lock()
221 defer mu.Unlock()
222 cp := *goal
223 return &cp
224 }
225 var persisted []*core.SessionGoal
226 persist := func(next *core.SessionGoal) error {
227 mu.Lock()
228 cp := *next
229 persisted = append(persisted, &cp)
230 goal = &cp
231 mu.Unlock()
232 return nil
233 }
234 client := &turnRecoveryModeClient{turns: [][]provider.Content{nil, nil}}
235 ag := core.NewAgent(client, "test-model", "", goalToolRegistry())
236 i := NewInteractive(InteractiveConfig{
237 Agent: ag,
238 CurrentGoal: current,
239 PersistGoal: persist,
240 PersistGoalRuntime: persist,
241 })
242 err := ag.Prompt(context.Background(), "work", nil, func(core.AgentEvent) {})
243 if !errors.Is(err, core.ErrIncompleteTurn) {
244 t.Fatalf("Prompt error = %v, want ErrIncompleteTurn", err)
245 }
246 // Exhaustion follows established blocking semantics: the goal is marked
247 // blocked with the incomplete-turn cause, never auto-cleared or rerun.
248 i.updateActiveGoal(core.GoalBlocked, "turn ended with an error: "+core.ErrIncompleteTurn.Error())
249 mu.Lock()
250 status := goal.Status
251 mu.Unlock()
252 if status != core.GoalBlocked {
253 t.Fatalf("goal status = %v, want blocked after exhaustion", status)
254 }
255 if len(persisted) == 0 {
256 t.Fatal("goal exhaustion was not persisted")
257 }
258}
259
260func TestTurnRecoveryEventJSONEncoding(t *testing.T) {
261 encoded := EventToJSON(core.EvTurnRecovery{Reason: core.TurnRecoveryReasonMissingAnswer, Attempt: 1})
262 raw, err := json.Marshal(encoded)
263 if err != nil {
264 t.Fatal(err)
265 }
266 var decoded map[string]any
267 if err := json.Unmarshal(raw, &decoded); err != nil {
268 t.Fatalf("turn_recovery event is not JSON: %v", err)
269 }
270 if decoded["type"] != "turn_recovery" {
271 t.Fatalf("type = %#v, want turn_recovery", decoded["type"])
272 }
273 if decoded["reason"] != core.TurnRecoveryReasonMissingAnswer {
274 t.Fatalf("reason = %#v, want missing_answer", decoded["reason"])
275 }
276 if decoded["attempt"] != float64(1) {
277 t.Fatalf("attempt = %#v, want 1", decoded["attempt"])
278 }
279}