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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "net/http"
   7         "os"
   8         "path/filepath"
   9         "strings"
  10         "testing"
  11         "time"
  12
  13         "github.com/stretchr/testify/require"
  14 )
  15
  16 func TestRebuildMergesDatesAndKeepsAgentFiles(t *testing.T) {
  17         agent := newRebuildTestAgent(t)
  18         old := []DatedMessage{
  19                 {
  20                         Kind: "note",
  21                         Text: "middle",
  22                         Date: "2025-03-02T12:00:00-03:00",
  23                 },
  24                 {
  25                         Kind: "user",
  26                         Text: "last",
  27                         Date: "2025-03-03T12:00:00Z",
  28                 },
  29         }
  30
  31         rebuildTestLog(t, agent.dir, old)
  32         require.NoError(t, os.MkdirAll(filepath.Join(agent.dir, "subagents", "BlueBird"), 0o700))
  33
  34         files := map[string]string{
  35                 "instructions.md":           "Keep exact answers.",
  36                 "schedules.json":            "[]\n",
  37                 "telegram-offset.json":      "123\n",
  38                 "work.json":                 "{\"running\":[]}\n",
  39                 "usage.jsonl":               "{\"model\":\"old\"}\n",
  40                 "subagents/BlueBird/report": "Finished.",
  41         }
  42
  43         for name, text := range files {
  44                 require.NoError(t, os.WriteFile(filepath.Join(agent.dir, name), []byte(text), 0o600))
  45         }
  46
  47         before := rebuildTestSnapshot(t, agent.dir)
  48         server := newModelFake(t)
  49         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
  50         var report strings.Builder
  51         added := []DatedMessage{
  52                 {
  53                         Kind: "work",
  54                         Text: strings.Repeat("Earlier message. ", 80),
  55                         Date: "2025-03-01T10:00:00-03:00",
  56                 },
  57                 {
  58                         Kind: "note",
  59                         Text: "equal date comes after old",
  60                         Date: old[0].Date,
  61                 },
  62         }
  63
  64         backup, err := rebuild(t.Context(), agent, added, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &report)
  65         require.NoError(t, err)
  66         require.Equal(t, before, rebuildTestSnapshot(t, backup))
  67
  68         entries, err := rebuildReadLog(t.Context(), agent.dir)
  69         require.NoError(t, err)
  70         require.Equal(t, []DatedMessage{added[0], old[0], added[1], old[1]}, entries)
  71
  72         for name, text := range files {
  73                 var data []byte
  74
  75                 data, err = os.ReadFile(filepath.Join(agent.dir, name))
  76                 require.NoError(t, err)
  77
  78                 if name == "usage.jsonl" {
  79                         require.True(t, strings.HasPrefix(string(data), text))
  80
  81                         continue
  82                 }
  83
  84                 require.Equal(t, text, string(data))
  85         }
  86
  87         require.Positive(t, server.Calls())
  88
  89         call := server.Call(0)
  90         require.Equal(t, agent.CacheKey(), call.body["prompt_cache_key"])
  91
  92         prompt, err := agent.Prompt()
  93         require.NoError(t, err)
  94         require.Contains(t, string(call.raw), strings.Trim(string(rebuildTestJSON(t, prompt)), "\""))
  95         require.Contains(t, report.String(), "cached (")
  96
  97         memory, err := openMemory(t.Context(), agent.dir, time.Now, newFakeSummarizer(), discardLog())
  98         require.NoError(t, err)
  99         require.NoError(t, memory.Settle(t.Context()))
 100         require.NoError(t, memory.Close())
 101 }
 102
 103 func TestRebuildFailureAndCancellationKeepOriginal(t *testing.T) {
 104         testCases := []struct {
 105                 name   string
 106                 cancel bool
 107         }{
 108                 {
 109                         name: "failed call keeps the original",
 110                 },
 111                 {
 112                         name:   "cancelled run keeps the original",
 113                         cancel: true,
 114                 },
 115         }
 116
 117         for _, tc := range testCases {
 118                 t.Run(tc.name, func(t *testing.T) {
 119                         agent := newRebuildTestAgent(t)
 120                         rebuildTestLog(t, agent.dir, []DatedMessage{{
 121                                 Kind: "user",
 122                                 Text: "Keep this.",
 123                                 Date: "2026-01-01T00:00:00Z",
 124                         }})
 125
 126                         before := rebuildTestSnapshot(t, agent.dir)
 127                         server := newModelFake(t)
 128                         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 129                                 w.WriteHeader(http.StatusBadRequest)
 130                         }
 131
 132                         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
 133                         ctx, cancel := context.WithCancel(t.Context())
 134                         defer cancel()
 135
 136                         if tc.cancel {
 137                                 cancel()
 138                         }
 139
 140                         backup, err := rebuild(ctx, agent, []DatedMessage{{
 141                                 Kind: "work",
 142                                 Text: strings.Repeat("long ", 200),
 143                                 Date: "2025-01-01T00:00:00Z",
 144                         }}, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &strings.Builder{})
 145                         require.Error(t, err)
 146                         require.Empty(t, backup)
 147                         require.Equal(t, before, rebuildTestSnapshot(t, agent.dir))
 148
 149                         pending, err := filepath.Glob(filepath.Join(filepath.Dir(agent.dir), ".agent-rebuild-*"))
 150                         require.NoError(t, err)
 151                         require.Empty(t, pending)
 152                 })
 153         }
 154 }
 155
 156 func TestRebuildKeepsOriginalWhenAnInvisibleMergeFails(t *testing.T) {
 157         agent := newRebuildTestAgent(t)
 158         old := []DatedMessage{
 159                 {
 160                         Kind: "user",
 161                         Text: strings.Repeat("a", 300),
 162                         Date: "2025-01-01T00:00:00Z",
 163                 },
 164                 {
 165                         Kind: "reply",
 166                         Text: strings.Repeat("b", 300),
 167                         Date: "2025-01-01T00:01:00Z",
 168                 },
 169         }
 170
 171         rebuildTestLog(t, agent.dir, old)
 172         before := rebuildTestSnapshot(t, agent.dir)
 173         server := newModelFake(t)
 174         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 175                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("")}, 100, 0, 5))
 176         }
 177
 178         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
 179         _, err := rebuild(t.Context(), agent, nil, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &strings.Builder{})
 180         require.ErrorContains(t, err, "summary tree has unfinished nodes")
 181         require.Equal(t, before, rebuildTestSnapshot(t, agent.dir))
 182 }
 183
 184 func TestRebuildRefusesUnfinishedSubagents(t *testing.T) {
 185         agent := newRebuildTestAgent(t)
 186         require.NoError(t, os.WriteFile(filepath.Join(agent.dir, "work.json"), []byte(`{"running":["BlueBird"]}`), 0o600))
 187
 188         require.ErrorContains(t, rebuildIdle(agent.dir), "unfinished work")
 189 }
 190
 191 func TestRebuildReadDoesNotRepairIncompleteLog(t *testing.T) {
 192         agent := newRebuildTestAgent(t)
 193         require.NoError(t, os.Mkdir(filepath.Join(agent.dir, "main"), 0o700))
 194
 195         path := filepath.Join(agent.dir, "main", "2026-01-01.jsonl")
 196         require.NoError(t, os.WriteFile(path, []byte("incomplete"), 0o600))
 197
 198         _, err := rebuildReadLog(t.Context(), agent.dir)
 199         require.ErrorContains(t, err, "incomplete line")
 200
 201         data, err := os.ReadFile(path)
 202         require.NoError(t, err)
 203         require.Equal(t, "incomplete", string(data))
 204 }
 205
 206 func TestRebuildReadStopsWhenTheContextEnds(t *testing.T) {
 207         agent := newRebuildTestAgent(t)
 208         rebuildTestLog(t, agent.dir, []DatedMessage{{
 209                 Kind: "note",
 210                 Text: "x",
 211                 Date: "2026-01-01T00:00:00Z",
 212         }})
 213
 214         ctx, cancel := context.WithCancel(t.Context())
 215         defer cancel()
 216
 217         cancel()
 218
 219         _, err := rebuildReadLog(ctx, agent.dir)
 220         require.ErrorIs(t, err, context.Canceled)
 221 }
 222
 223 func rebuildTestJSON(t *testing.T, value string) []byte {
 224         t.Helper()
 225
 226         data, err := json.Marshal(value)
 227         require.NoError(t, err)
 228
 229         return data
 230 }