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