1package modes
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/json"
8 "errors"
9 "sync"
10 "testing"
11
12 "github.com/bnema/zut/packages/core"
13 "github.com/bnema/zut/packages/provider"
14)
15
16type hostRecoveryClient struct {
17 mu sync.Mutex
18 calls int
19 retried chan provider.Request
20 retryOverflow bool
21}
22
23func (c *hostRecoveryClient) Name() string { return "host-context-recovery-test" }
24
25func (c *hostRecoveryClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
26 c.mu.Lock()
27 c.calls++
28 call := c.calls
29 c.mu.Unlock()
30
31 events := make(chan provider.Event, 2)
32 go func() {
33 defer close(events)
34 switch call {
35 case 1:
36 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("provider error: input exceeds the context window")}
37 case 2:
38 events <- provider.EventTextDelta{Delta: "summary"}
39 events <- provider.EventDone{Stop: provider.StopEnd}
40 default:
41 c.retried <- req
42 if c.retryOverflow {
43 events <- provider.EventDone{Stop: provider.StopError, Err: errors.New("context window exceeded after compaction")}
44 return
45 }
46 events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
47 Role: provider.RoleAssistant,
48 Content: []provider.Content{provider.TextBlock{Text: "recovered"}},
49 }}
50 }
51 }()
52 return events, nil
53}
54
55func TestPromptWithContextRecoveryCompactsOnceAndSuppressesInitialOverflow(t *testing.T) {
56 client := &hostRecoveryClient{retried: make(chan provider.Request, 1)}
57 agent := core.NewAgent(client, "test-model", "", nil)
58 agent.SetMessages([]provider.Message{
59 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
60 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
61 })
62
63 var sawOverflowTurnEnd bool
64 compacted := false
65 result, err := PromptWithContextRecovery(context.Background(), agent, "finish the task", nil, func(event core.AgentEvent) {
66 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
67 sawOverflowTurnEnd = true
68 }
69 }, ContextRecoveryOptions{
70 PersistCompaction: func([]provider.Message) error {
71 compacted = true
72 return nil
73 },
74 SuppressInitialOverflowEvent: true,
75 })
76 if err != nil {
77 t.Fatalf("PromptWithContextRecovery returned %v", err)
78 }
79 if !result.Compacted || !compacted {
80 t.Fatalf("compaction result/persistence = %v/%v, want true/true", result.Compacted, compacted)
81 }
82 if sawOverflowTurnEnd {
83 t.Fatal("recoverable initial overflow turn_end reached the sink")
84 }
85 if got := recoveredAssistantText(agent.Messages()[result.OutputStart:]); got != "recovered" {
86 t.Fatalf("recovery output = %q, want recovered", got)
87 }
88 request := <-client.retried
89 if got := requestUserTextCount(request, "finish the task"); got != 1 {
90 t.Fatalf("retried request prompt count = %d, want 1: %#v", got, request.Messages)
91 }
92}
93
94func recoveredAssistantText(messages []provider.Message) string {
95 for _, message := range messages {
96 if message.Role != provider.RoleAssistant {
97 continue
98 }
99 for _, content := range message.Content {
100 if text, ok := content.(provider.TextBlock); ok {
101 return text.Text
102 }
103 }
104 }
105 return ""
106}
107
108func TestPromptWithContextRecoveryReportsRetryOverflowToSink(t *testing.T) {
109 client := &hostRecoveryClient{retried: make(chan provider.Request, 1), retryOverflow: true}
110 agent := core.NewAgent(client, "test-model", "", nil)
111 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
112
113 overflowTurnEnds := 0
114 _, err := PromptWithContextRecovery(context.Background(), agent, "finish the task", nil, func(event core.AgentEvent) {
115 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
116 overflowTurnEnds++
117 }
118 }, ContextRecoveryOptions{SuppressInitialOverflowEvent: true})
119 if !provider.IsContextOverflowError(err) {
120 t.Fatalf("PromptWithContextRecovery error = %v, want retry context overflow", err)
121 }
122 if overflowTurnEnds != 1 {
123 t.Fatalf("overflow turn_end count = %d, want retry overflow only", overflowTurnEnds)
124 }
125 client.mu.Lock()
126 calls := client.calls
127 client.mu.Unlock()
128 if calls != 3 {
129 t.Fatalf("provider calls = %d, want prompt, compaction, and one continuation", calls)
130 }
131}
132
133func newRecoveredModeAgent() *core.Agent {
134 return core.NewAgent(&hostRecoveryClient{retried: make(chan provider.Request, 1)}, "test-model", "", nil)
135}
136
137func TestRunPrintRecoversContextOverflowWithoutLeakingIntermediateOutput(t *testing.T) {
138 agent := newRecoveredModeAgent()
139 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
140 var output bytes.Buffer
141
142 if _, err := RunPrint(context.Background(), agent, "finish the task", nil, &output); err != nil {
143 t.Fatalf("RunPrint returned %v", err)
144 }
145 if got := output.String(); got != "recovered\n" {
146 t.Fatalf("stdout = %q, want recovered only", got)
147 }
148}
149
150func TestRunStreamRecoversContextOverflowWithoutDiagnosticsLeak(t *testing.T) {
151 agent := newRecoveredModeAgent()
152 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
153 var output, diagnostics bytes.Buffer
154
155 if err := RunStreamWithDiag(context.Background(), agent, "finish the task", nil, &output, &diagnostics); err != nil {
156 t.Fatalf("RunStreamWithDiag returned %v", err)
157 }
158 if got := output.String(); got != "recovered\n" {
159 t.Fatalf("stdout = %q, want recovered assistant text once", got)
160 }
161 if got := diagnostics.String(); got != "" {
162 t.Fatalf("stderr diagnostics = %q, want no context-overflow diagnostic", got)
163 }
164}
165
166func TestRunJSONRecoversContextOverflowWithJSONLOnly(t *testing.T) {
167 agent := newRecoveredModeAgent()
168 agent.SetMessages([]provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "prior context"}}}})
169 var output bytes.Buffer
170
171 if err := RunJSON(context.Background(), agent, "finish the task", nil, &output); err != nil {
172 t.Fatalf("RunJSON returned %v", err)
173 }
174
175 var (
176 turnEnds int
177 done int
178 )
179 scanner := bufio.NewScanner(&output)
180 for scanner.Scan() {
181 var event map[string]any
182 if err := json.Unmarshal(scanner.Bytes(), &event); err != nil {
183 t.Fatalf("stdout line is not JSONL: %q: %v", scanner.Text(), err)
184 }
185 switch event["type"] {
186 case "error":
187 t.Fatalf("recoverable overflow wrote error frame: %#v", event)
188 case "turn_end":
189 turnEnds++
190 if _, hasError := event["error"]; hasError {
191 t.Fatalf("recoverable overflow wrote terminal turn_end: %#v", event)
192 }
193 case "done":
194 done++
195 }
196 }
197 if err := scanner.Err(); err != nil {
198 t.Fatalf("scan JSONL: %v", err)
199 }
200 if turnEnds != 1 || done != 1 {
201 t.Fatalf("turn_end/done count = %d/%d, want one recovered terminal sequence", turnEnds, done)
202 }
203}