1package modes
2
3import (
4 "context"
5 "encoding/json"
6 "strings"
7 "testing"
8
9 "github.com/bnema/zut/packages/core"
10 "github.com/bnema/zut/packages/provider"
11 "github.com/bnema/zut/packages/tui"
12)
13
14// scriptedTurnsClient replays one scripted event list per provider request.
15type scriptedTurnsClient struct {
16 turns [][]provider.Event
17 call int
18}
19
20func (c *scriptedTurnsClient) Name() string { return "scripted" }
21
22func (c *scriptedTurnsClient) Stream(ctx context.Context, _ provider.Request) (<-chan provider.Event, error) {
23 events := c.turns[min(c.call, len(c.turns)-1)]
24 c.call++
25 out := make(chan provider.Event, len(events))
26 for _, ev := range events {
27 out <- ev
28 }
29 close(out)
30 return out, nil
31}
32
33type echoTool struct{}
34
35func (echoTool) Name() string { return "echo" }
36func (echoTool) Description() string { return "echo" }
37func (echoTool) Schema() json.RawMessage {
38 return json.RawMessage(`{"type":"object"}`)
39}
40func (echoTool) Execute(context.Context, json.RawMessage, func(string)) (core.ToolResult, error) {
41 return core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: "tool output"}}}, nil
42}
43
44func textTurn(text string, stop provider.StopReason, extra ...provider.Content) []provider.Event {
45 content := append([]provider.Content{provider.TextBlock{Text: text}}, extra...)
46 var events []provider.Event
47 // Fat chunks like the Anthropic OAuth path: the provider is far ahead of
48 // the pacer when the message completes.
49 for chunk := range strings.SplitSeq(text, " ") {
50 events = append(events, provider.EventTextDelta{Delta: chunk + " "})
51 }
52 return append(events, provider.EventDone{Stop: stop, Message: provider.Message{Role: provider.RoleAssistant, Content: content}})
53}
54
55// countVisible reports how many times text appears in the visible chat.
56func countVisible(i *Interactive, text string) int {
57 i.mu.Lock()
58 defer i.mu.Unlock()
59 return strings.Count(stripANSIBytes(strings.Join(i.buildChatLocked(120), "\n")), text)
60}
61
62// newStreamRaceInteractive wires an agent to an interactive view without
63// the pacer goroutine, so the test controls exactly when text is painted.
64func newStreamRaceInteractive(t *testing.T, client provider.Client, tools core.Registry) (*Interactive, *core.Agent) {
65 t.Helper()
66 ag := core.NewAgent(client, "test-model", "", tools)
67 i := &Interactive{
68 agent: ag,
69 view: &tui.View{Theme: tui.Dark},
70 toolCalls: map[string]*tui.ToolCallView{},
71 dirty: make(chan struct{}, 1),
72 busy: true,
73 }
74 return i, ag
75}
76
77// Race 1: the agent appends the finished reply and persists it before the UI
78// receives EvAssistantMessage. A frame drawn in that window must not show the
79// reply both as a transcript message and as partial live text.
80func TestStreamReplyNotDuplicatedWhileMessageIsPersisted(t *testing.T) {
81 const reply = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu"
82 client := &scriptedTurnsClient{turns: [][]provider.Event{textTurn(reply, provider.StopEnd)}}
83 i, ag := newStreamRaceInteractive(t, client, nil)
84
85 var seen []int
86 ag.OnMessageAppended = func(m provider.Message) {
87 if m.Role != provider.RoleAssistant {
88 return
89 }
90 // Paint part of the reply, as the pacer would have by now.
91 i.mu.Lock()
92 drainTicks(&i.stream, 3)
93 i.mu.Unlock()
94 seen = append(seen, countVisible(i, "alpha"))
95 }
96 if err := ag.Prompt(context.Background(), "question", nil, i.handleEvent); err != nil {
97 t.Fatal(err)
98 }
99 if len(seen) != 1 || seen[0] != 1 {
100 t.Fatalf("reply painted %v times while the message was being persisted, want exactly once", seen)
101 }
102}
103
104// Race 2: after a text+tool reply, the tool result lands in the transcript
105// while the pacer is still painting the text. The reply must stay single.
106func TestStreamReplyNotDuplicatedWhenToolResultArrivesEarly(t *testing.T) {
107 const reply = "let me check that file before answering the question properly"
108 toolCall := provider.ToolCallBlock{ID: "call-1", Name: "echo", Arguments: json.RawMessage(`{}`)}
109 client := &scriptedTurnsClient{turns: [][]provider.Event{
110 textTurn(reply, provider.StopToolUse, toolCall),
111 textTurn("done", provider.StopEnd),
112 }}
113 i, ag := newStreamRaceInteractive(t, client, core.Registry{"echo": echoTool{}})
114
115 var seen []int
116 sawLive := false
117 ag.OnMessageAppended = func(m provider.Message) {
118 if m.Role != provider.RoleTool {
119 return
120 }
121 // Tool result appended; the pacer has painted only part of the reply.
122 i.mu.Lock()
123 drainTicks(&i.stream, 3)
124 sawLive = sawLive || i.stream.Active()
125 i.mu.Unlock()
126 seen = append(seen, countVisible(i, "let me"))
127 }
128 if err := ag.Prompt(context.Background(), "question", nil, i.handleEvent); err != nil {
129 t.Fatal(err)
130 }
131 if !sawLive {
132 t.Fatal("test setup: pacer finished before the tool result was appended")
133 }
134 if len(seen) != 1 || seen[0] != 1 {
135 t.Fatalf("reply painted %v times after the tool result arrived, want exactly once", seen)
136 }
137}
138
139// Esc usually surfaces as a context error, not StopAborted. Buffered text
140// must stop at once instead of typing on and then vanishing.
141func TestStreamStopsImmediatelyOnCancel(t *testing.T) {
142 i := &Interactive{view: &tui.View{Theme: tui.Dark}, toolCalls: map[string]*tui.ToolCallView{}}
143 i.handleEvent(core.EvAssistantStart{})
144 i.handleEvent(core.EvTextDelta{Delta: strings.Repeat("x", 500)})
145 i.handleEvent(core.EvTurnEnd{Stop: provider.StopError, Err: context.Canceled})
146 if i.stream.Active() || i.stream.Text() != "" {
147 t.Fatalf("stream still active after cancel: active=%t text=%d", i.stream.Active(), len(i.stream.Text()))
148 }
149}
150
151// Replacing the agent (/model, /login) while text drains must drop the old
152// agent's transcript anchor.
153func TestStreamResetOnAgentReplacement(t *testing.T) {
154 i := &Interactive{view: &tui.View{Theme: tui.Dark}, toolCalls: map[string]*tui.ToolCallView{}}
155 i.stream.Start(3, 7)
156 i.stream.Push("draining")
157 i.stream.Finish()
158 i.prepareReplacementAgentLocked(nil)
159 if i.stream.Active() {
160 t.Fatal("stream anchor survived an agent replacement")
161 }
162}
163
164// After the pacer drains, the transcript shows the reply exactly once and
165// the live block is gone.
166func TestStreamRevealsTranscriptWhenPacerDrains(t *testing.T) {
167 const reply = "one two three four five six seven eight nine ten"
168 client := &scriptedTurnsClient{turns: [][]provider.Event{textTurn(reply, provider.StopEnd)}}
169 i, ag := newStreamRaceInteractive(t, client, nil)
170 if err := ag.Prompt(context.Background(), "question", nil, i.handleEvent); err != nil {
171 t.Fatal(err)
172 }
173
174 i.mu.Lock()
175 drainTicks(&i.stream, 3)
176 active := i.stream.Active()
177 i.mu.Unlock()
178 if !active {
179 t.Fatal("test setup: stream drained after one tick")
180 }
181 if got := countVisible(i, "one two"); got != 1 {
182 t.Fatalf("reply painted %d times before draining, want 1", got)
183 }
184
185 i.mu.Lock()
186 for i.stream.Active() {
187 i.stream.Tick()
188 }
189 i.mu.Unlock()
190 if got := countVisible(i, reply); got != 1 {
191 t.Fatalf("reply painted %d times after draining, want 1", got)
192 }
193}
194
195// A backlog far larger than the minimum rate is painted within a bounded
196// number of ticks, so a fast provider never leaves the UI trailing far behind.
197func TestStreamPacerCatchesUpWithLargeBacklog(t *testing.T) {
198 var s streamPresenter
199 s.Start(-1, 0)
200 s.Push(strings.Repeat("x", 20000))
201 s.Finish()
202 ticks := 0
203 for s.Active() {
204 s.Tick()
205 ticks++
206 }
207 if limit := paintCatchUpTicks * 8; ticks > limit {
208 t.Fatalf("drained a 20k-rune backlog in %d ticks, want <= %d", ticks, limit)
209 }
210}
211
212// A transcript replaced by a shorter one (compaction, /clear) must never be
213// clipped by an anchor that described the old transcript.
214func TestStreamAnchorIgnoredAfterTranscriptShrinks(t *testing.T) {
215 var s streamPresenter
216 s.Start(10, 42)
217 s.Push("text")
218 if _, _, ok := s.transcriptLimit(3); ok {
219 t.Fatal("anchor clipped a transcript shorter than itself")
220 }
221 if limit, rev, ok := s.transcriptLimit(12); !ok || limit != 10 || rev != 42 {
222 t.Fatalf("transcriptLimit(12) = %d, %d, %t; want 10, 42, true", limit, rev, ok)
223 }
224}