]> git.giacomolli.uk - onechat/blob - turn_test.go
Import WhatsApp history into a child
[onechat] / turn_test.go
1 package main
2
3 import (
4         "context"
5         "encoding/json"
6         "net/http"
7         "path/filepath"
8         "strings"
9         "sync"
10         "testing"
11         "time"
12
13         "github.com/stretchr/testify/require"
14 )
15
16 func TestTurnRunsToolsUntilTheFinalAnswer(t *testing.T) {
17         server := newModelFake(t)
18         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
19                 if call.number == 0 {
20                         modelWriteEvents(w, modelCompletedEvent([]any{
21                                 modelReasoningItem("rs_1", "opaque thought"),
22                                 modelFunctionCallItem("fc_1", "call_1", "shell", `{"command":"echo hi"}`),
23                         }, 100, 0, 5))
24
25                         return
26                 }
27
28                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("It said hi.")}, 120, 100, 5))
29         }
30
31         test := newTestTurnLoop(t, server)
32         _, err := test.loop.Submit(terminalSource, "Run echo hi.")
33         require.NoError(t, err)
34         test.events.WaitFor(t, "idle", 1)
35         require.Equal(t, 2, server.Calls())
36
37         first := server.Call(0).body
38         second := server.Call(1).body
39         instructions, err := test.agent.Prompt()
40         require.NoError(t, err)
41         require.Equal(t, instructions, first["instructions"])
42         require.Equal(t, test.agent.config.Model, first["model"])
43
44         // The second request begins with the whole first one.
45         firstInput := first["input"].([]any)
46         secondInput := second["input"].([]any)
47         require.Equal(t, firstInput, secondInput[:len(firstInput)])
48         require.Equal(t, []any{
49                 map[string]any{"type": "reasoning", "encrypted_content": "opaque thought", "summary": []any{}},
50                 map[string]any{"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": `{"command":"echo hi"}`},
51                 map[string]any{"type": "function_call_output", "call_id": "call_1", "output": "hi\n"},
52         }, secondInput[len(firstInput):])
53
54         require.Equal(t, []string{
55                 "user: Run echo hi.",
56                 `tool: shell {"command":"echo hi"}`,
57                 "echo: hi\n",
58                 "reply: It said hi.",
59         }, test.Log(t))
60         require.Equal(t, []string{"user", "busy", "tool", "echo", "reply", "idle"}, test.events.Kinds())
61
62         var turns []string
63
64         for _, record := range modelReadUsage(t, test.usagePath) {
65                 turns = append(turns, record.Turn)
66         }
67
68         require.Equal(t, []string{"first", "next"}, turns)
69 }
70
71 func TestFirstRequestHoldsTheViewTheStateAndTheMessage(t *testing.T) {
72         server := newModelFake(t)
73         test := newTestTurnLoop(t, server)
74         _, err := test.memory.Append("user", "an earlier message")
75         require.NoError(t, err)
76
77         _, err = test.loop.Submit(terminalSource, "The new message.")
78         require.NoError(t, err)
79         test.events.WaitFor(t, "idle", 1)
80
81         require.Equal(t, []string{
82                 "<chat>\n0+1|user: an earlier message\n",
83                 "</chat>\n\nNow: Thursday 2026-01-01 12:00 +00:00",
84                 "The new message.",
85         }, turnInputTexts(t, server.Call(0)))
86 }
87
88 func TestMessageDuringATurnJoinsTheNextRequest(t *testing.T) {
89         server := newModelFake(t)
90         arrived := make(chan struct{})
91         release := make(chan struct{})
92         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
93                 if call.number == 0 {
94                         close(arrived)
95                         <-release
96                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "date", `{"id":0}`)}, 100, 0, 5))
97
98                         return
99                 }
100
101                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("done")}, 100, 0, 5))
102         }
103
104         test := newTestTurnLoop(t, server)
105         _, err := test.loop.Submit(terminalSource, "first")
106         require.NoError(t, err)
107
108         <-arrived
109
110         _, err = test.loop.Submit(terminalSource, "second, during the turn")
111         require.NoError(t, err)
112         close(release)
113
114         test.events.WaitFor(t, "idle", 1)
115
116         input := server.Call(1).body["input"].([]any)
117         require.Equal(t, map[string]any{
118                 "role":    "user",
119                 "content": []any{map[string]any{"type": "input_text", "text": "second, during the turn"}},
120         }, input[len(input)-1])
121
122         time.Sleep(50 * time.Millisecond)
123         require.Equal(t, 2, server.Calls())
124         require.Equal(t, 1, test.events.Count("busy"))
125 }
126
127 func TestMessageDuringTheFinalRequestStartsANewTurn(t *testing.T) {
128         server := newModelFake(t)
129         arrived := make(chan struct{})
130         release := make(chan struct{})
131         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
132                 if call.number == 0 {
133                         close(arrived)
134                         <-release
135                 }
136
137                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("answer")}, 100, 0, 5))
138         }
139
140         test := newTestTurnLoop(t, server)
141         _, err := test.loop.Submit(terminalSource, "first")
142         require.NoError(t, err)
143
144         <-arrived
145
146         _, err = test.loop.Submit(terminalSource, "second")
147         require.NoError(t, err)
148         close(release)
149
150         test.events.WaitFor(t, "idle", 2)
151         require.Equal(t, 2, server.Calls())
152
153         texts := turnInputTexts(t, server.Call(1))
154         require.Equal(t, "second", texts[len(texts)-1])
155         require.Contains(t, strings.Join(texts, ""), "|reply: answer\n")
156 }
157
158 func TestTurnWaitsForTheSummaries(t *testing.T) {
159         server := newModelFake(t)
160         summarizer := newFakeSummarizer()
161         summarizer.gate = make(chan struct{})
162         test := newTestTurnLoopWith(t, server, summarizer)
163
164         _, err := test.memory.Append("echo", strings.Repeat("long output ", 100))
165         require.NoError(t, err)
166
167         _, err = test.loop.Submit(terminalSource, "go on")
168         require.NoError(t, err)
169
170         time.Sleep(50 * time.Millisecond)
171         require.Equal(t, 0, server.Calls())
172
173         close(summarizer.gate)
174         test.events.WaitFor(t, "idle", 1)
175         require.Equal(t, 1, server.Calls())
176         require.NotContains(t, turnInputTexts(t, server.Call(0))[0], viewPending)
177 }
178
179 func TestInterruptEndsTheTurn(t *testing.T) {
180         server := newModelFake(t)
181         arrived := make(chan struct{})
182         release := make(chan struct{})
183         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
184                 if call.number == 0 {
185                         close(arrived)
186                         <-release
187                 }
188
189                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("answer")}, 100, 0, 5))
190         }
191
192         test := newTestTurnLoop(t, server)
193         _, err := test.loop.Submit(terminalSource, "a long task")
194         require.NoError(t, err)
195
196         <-arrived
197         require.True(t, test.loop.Interrupt())
198         test.events.WaitFor(t, "idle", 1)
199         close(release)
200
201         require.Equal(t, "interrupted", test.events.Last("error"))
202         require.Equal(t, []string{"user: a long task"}, test.Log(t))
203         require.False(t, test.loop.Interrupt())
204
205         _, err = test.loop.Submit(terminalSource, "again")
206         require.NoError(t, err)
207         test.events.WaitFor(t, "idle", 2)
208         require.Equal(t, "reply: answer", test.Log(t)[2])
209         require.Empty(t, test.logs.String())
210 }
211
212 func TestFailedCallEndsTheTurnWithoutARetry(t *testing.T) {
213         server := newModelFake(t)
214         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
215                 http.Error(w, "busy", http.StatusServiceUnavailable)
216         }
217
218         test := newTestTurnLoop(t, server)
219         _, err := test.loop.Submit(terminalSource, "hello")
220         require.NoError(t, err)
221         test.events.WaitFor(t, "idle", 1)
222
223         time.Sleep(50 * time.Millisecond)
224         require.Equal(t, 1, server.Calls())
225         require.Contains(t, test.events.Last("error"), "HTTP 503")
226         require.Equal(t, 1, strings.Count(test.logs.String(), "\n"))
227         require.Contains(t, test.logs.String(), "<3>turn failed: ")
228         require.Contains(t, test.logs.String(), "HTTP 503")
229         require.NotContains(t, test.logs.String(), "hello")
230 }
231
232 func TestViewMessagesOnlyGrowBetweenTurns(t *testing.T) {
233         server := newModelFake(t)
234         test := newTestTurnLoop(t, server)
235         for n, text := range []string{"one", "two", "three"} {
236                 _, err := test.loop.Submit(terminalSource, text)
237                 require.NoError(t, err)
238                 test.events.WaitFor(t, "idle", n+1)
239         }
240
241         first := server.Call(1).body["input"].([]any)
242         second := server.Call(2).body["input"].([]any)
243
244         require.Len(t, first, 4)
245         require.Len(t, second, 5)
246         require.Equal(t, first[:2], second[:2])
247         require.Equal(t, []string{
248                 "<chat>\n",
249                 "0+1|user: one\n1+1|reply: a line\n",
250                 "2+1|user: two\n3+1|reply: a line\n",
251                 "</chat>\n\nNow: Thursday 2026-01-01 12:00 +00:00",
252                 "three",
253         }, turnInputTexts(t, server.Call(2)))
254 }
255
256 // TurnTestLoop is a loop that runs in the background of a test, with its parts.
257 type TurnTestLoop struct {
258         agent     *Agent
259         memory    *Memory
260         loop      *TurnLoop
261         events    *TurnEventRecorder
262         logs      *SyncBuffer
263         usagePath string
264 }
265
266 // Log returns the logged messages as "kind: text".
267 func (l *TurnTestLoop) Log(t *testing.T) []string {
268         t.Helper()
269
270         var lines []string
271
272         for id := range l.memory.Len() {
273                 message, err := l.memory.Message(id)
274                 require.NoError(t, err)
275
276                 lines = append(lines, message.Kind+": "+message.Text)
277         }
278
279         return lines
280 }
281
282 // TurnEventRecorder keeps the events of a loop.
283 type TurnEventRecorder struct {
284         mutex  sync.Mutex
285         events []Event
286 }
287
288 // WaitFor waits until count events of kind came.
289 func (r *TurnEventRecorder) WaitFor(t *testing.T, kind string, count int) {
290         t.Helper()
291
292         require.Eventually(t, func() bool {
293                 return r.Count(kind) >= count
294         }, 10*time.Second, time.Millisecond)
295 }
296
297 func (r *TurnEventRecorder) Kinds() []string {
298         r.mutex.Lock()
299         defer r.mutex.Unlock()
300
301         kinds := make([]string, len(r.events))
302         for n, event := range r.events {
303                 kinds[n] = event.Kind
304         }
305
306         return kinds
307 }
308
309 func (r *TurnEventRecorder) Count(kind string) int {
310         count := 0
311         for _, each := range r.Kinds() {
312                 if each == kind {
313                         count++
314                 }
315         }
316
317         return count
318 }
319
320 // Last returns the text of the last event of kind.
321 func (r *TurnEventRecorder) Last(kind string) string {
322         return r.LastEvent(kind).Text
323 }
324
325 // LastEvent returns the last event of kind.
326 func (r *TurnEventRecorder) LastEvent(kind string) Event {
327         r.mutex.Lock()
328         defer r.mutex.Unlock()
329
330         for n := len(r.events) - 1; n >= 0; n-- {
331                 if r.events[n].Kind == kind {
332                         return r.events[n]
333                 }
334         }
335
336         return Event{}
337 }
338
339 func (r *TurnEventRecorder) Add(event Event) {
340         r.mutex.Lock()
341         defer r.mutex.Unlock()
342
343         r.events = append(r.events, event)
344 }
345
346 func newTestTurnLoop(t *testing.T, server *ModelFake) *TurnTestLoop {
347         t.Helper()
348
349         return newTestTurnLoopWith(t, server, newFakeSummarizer())
350 }
351
352 // turnInputTexts returns the text of each user message of the input of call.
353 func turnInputTexts(t *testing.T, call ModelFakeCall) []string {
354         t.Helper()
355
356         var body struct {
357                 Input []struct {
358                         Role    string `json:"role"`
359                         Content []struct {
360                                 Text string `json:"text"`
361                         } `json:"content"`
362                 } `json:"input"`
363         }
364
365         require.NoError(t, json.Unmarshal(call.raw, &body))
366
367         var texts []string
368
369         for _, item := range body.Input {
370                 if item.Role != "user" {
371                         continue
372                 }
373
374                 var text strings.Builder
375
376                 for _, part := range item.Content {
377                         text.WriteString(part.Text)
378                 }
379
380                 texts = append(texts, text.String())
381         }
382
383         return texts
384 }
385
386 // newTestTurnLoopWith returns a running loop on a new agent whose builds use
387 // summarizer and whose commands run in a temporary directory.
388 func newTestTurnLoopWith(t *testing.T, server *ModelFake, summarizer Summarizer) *TurnTestLoop {
389         t.Helper()
390
391         agent, err := openAgent(t.TempDir())
392         require.NoError(t, err)
393
394         agent.config.Workdir = t.TempDir()
395         clock := newTestClock()
396         memory, err := openMemory(t.Context(), t.TempDir(), clock.Now, summarizer, discardLog())
397         require.NoError(t, err)
398
399         test := &TurnTestLoop{
400                 agent:     agent,
401                 memory:    memory,
402                 events:    &TurnEventRecorder{},
403                 logs:      &SyncBuffer{},
404                 usagePath: filepath.Join(t.TempDir(), "usage.jsonl"),
405         }
406
407         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), test.usagePath)
408         whatsapp, _, _ := whatsappStartTest(t, t.TempDir(), clock)
409         test.loop, err = newTurnLoop(agent, memory, newTestSubagents(t), whatsapp, newTestSchedules(t, clock.Now), client, clock.Now, daemonTestLogger(test.logs), test.events.Add)
410         require.NoError(t, err)
411
412         ctx, cancel := context.WithCancel(t.Context())
413         done := make(chan error)
414         go func() {
415                 done <- test.loop.Run(ctx)
416         }()
417
418         t.Cleanup(func() {
419                 cancel()
420                 require.ErrorIs(t, <-done, context.Canceled)
421                 require.NoError(t, memory.Close())
422         })
423
424         return test
425 }