git.giacomolli.uk - onechat/blob - turn_test.go

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