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

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