git.giacomolli.uk - onechat/blob - 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 // TurnEventRecorder keeps the events of a loop.
 267 type TurnEventRecorder struct {
 268         mutex  sync.Mutex
 269         events []Event
 270 }
 271
 272 func newTestTurnLoop(t *testing.T, server *ModelFake) *TurnTestLoop {
 273         t.Helper()
 274
 275         return newTestTurnLoopWith(t, server, newFakeSummarizer())
 276 }
 277
 278 // newTestTurnLoopWith returns a running loop on a new agent whose builds use
 279 // summarizer and whose commands run in a temporary directory.
 280 func newTestTurnLoopWith(t *testing.T, server *ModelFake, summarizer Summarizer) *TurnTestLoop {
 281         t.Helper()
 282
 283         agent, err := openAgent(filepath.Join(t.TempDir(), "agent"))
 284         require.NoError(t, err)
 285
 286         agent.config.Workdir = t.TempDir()
 287         clock := newTestClock()
 288         memory, err := openMemory(t.Context(), t.TempDir(), clock.Now, summarizer, discardLog())
 289         require.NoError(t, err)
 290
 291         test := &TurnTestLoop{
 292                 agent:     agent,
 293                 memory:    memory,
 294                 events:    &TurnEventRecorder{},
 295                 logs:      &SyncBuffer{},
 296                 usagePath: filepath.Join(t.TempDir(), "usage.jsonl"),
 297         }
 298
 299         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), test.usagePath)
 300         whatsapp, _, _ := whatsappStartTest(t, t.TempDir(), clock)
 301         test.loop, err = newTurnLoop(agent, memory, newTestSubagents(t), whatsapp, newTestSchedules(t, clock.Now), client, clock.Now, daemonTestLogger(test.logs), test.events.Add)
 302         require.NoError(t, err)
 303
 304         ctx, cancel := context.WithCancel(t.Context())
 305         done := make(chan error)
 306         go func() {
 307                 done <- test.loop.Run(ctx)
 308         }()
 309
 310         t.Cleanup(func() {
 311                 cancel()
 312                 require.ErrorIs(t, <-done, context.Canceled)
 313                 require.NoError(t, memory.Close())
 314         })
 315
 316         return test
 317 }
 318
 319 // Log returns the logged messages as "kind: text".
 320 func (l *TurnTestLoop) Log(t *testing.T) []string {
 321         t.Helper()
 322
 323         var lines []string
 324
 325         for id := range l.memory.Len() {
 326                 message, err := l.memory.Message(id)
 327                 require.NoError(t, err)
 328
 329                 lines = append(lines, message.Kind+": "+message.Text)
 330         }
 331
 332         return lines
 333 }
 334
 335 // WaitFor waits until count events of kind came.
 336 func (r *TurnEventRecorder) WaitFor(t *testing.T, kind string, count int) {
 337         t.Helper()
 338
 339         require.Eventually(t, func() bool {
 340                 return r.Count(kind) >= count
 341         }, 10*time.Second, time.Millisecond)
 342 }
 343
 344 func (r *TurnEventRecorder) Kinds() []string {
 345         r.mutex.Lock()
 346         defer r.mutex.Unlock()
 347
 348         kinds := make([]string, len(r.events))
 349         for n, event := range r.events {
 350                 kinds[n] = event.Kind
 351         }
 352
 353         return kinds
 354 }
 355
 356 func (r *TurnEventRecorder) Count(kind string) int {
 357         count := 0
 358         for _, each := range r.Kinds() {
 359                 if each == kind {
 360                         count++
 361                 }
 362         }
 363
 364         return count
 365 }
 366
 367 // Last returns the text of the last event of kind.
 368 func (r *TurnEventRecorder) Last(kind string) string {
 369         return r.LastEvent(kind).Text
 370 }
 371
 372 // LastEvent returns the last event of kind.
 373 func (r *TurnEventRecorder) LastEvent(kind string) Event {
 374         r.mutex.Lock()
 375         defer r.mutex.Unlock()
 376
 377         for n := len(r.events) - 1; n >= 0; n-- {
 378                 if r.events[n].Kind == kind {
 379                         return r.events[n]
 380                 }
 381         }
 382
 383         return Event{}
 384 }
 385
 386 func (r *TurnEventRecorder) Add(event Event) {
 387         r.mutex.Lock()
 388         defer r.mutex.Unlock()
 389
 390         r.events = append(r.events, event)
 391 }
 392
 393 // turnInputTexts returns the text of each user message of the input of call.
 394 func turnInputTexts(t *testing.T, call ModelFakeCall) []string {
 395         t.Helper()
 396
 397         var body struct {
 398                 Input []struct {
 399                         Role    string `json:"role"`
 400                         Content []struct {
 401                                 Text string `json:"text"`
 402                         } `json:"content"`
 403                 } `json:"input"`
 404         }
 405
 406         require.NoError(t, json.Unmarshal(call.raw, &body))
 407
 408         var texts []string
 409
 410         for _, item := range body.Input {
 411                 if item.Role != "user" {
 412                         continue
 413                 }
 414
 415                 var text strings.Builder
 416
 417                 for _, part := range item.Content {
 418                         text.WriteString(part.Text)
 419                 }
 420
 421                 texts = append(texts, text.String())
 422         }
 423
 424         return texts
 425 }
 426
 427 func modelReasoningItem(id, encrypted string) map[string]any {
 428         return map[string]any{
 429                 "type":              "reasoning",
 430                 "id":                id,
 431                 "encrypted_content": encrypted,
 432                 "summary":           []any{},
 433         }
 434 }