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

   1 package main
   2
   3 import (
   4         "archive/tar"
   5         "context"
   6         "database/sql"
   7         "io"
   8         "os"
   9         "path/filepath"
  10         "strings"
  11         "sync"
  12         "testing"
  13         "time"
  14
  15         "github.com/stretchr/testify/require"
  16 )
  17
  18 func TestRunningDaemonBackupRestoresTheChatAndLinks(t *testing.T) {
  19         dir := t.TempDir()
  20         whatsappTestSession(t, dir, true)
  21
  22         files := map[string]string{
  23                 "auth.json":                  `{"access_token":"fake-secret"}`,
  24                 "telegram.json":              `{"token":"fake-token","owner":123}`,
  25                 "config.json":                `{"endpoint":"api"}`,
  26                 "agent/instructions.md":      "Keep the instructions.\n",
  27                 "agent/telegram-offset.json": "23\n",
  28         }
  29
  30         for name, data := range files {
  31                 path := filepath.Join(dir, name)
  32                 require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
  33                 require.NoError(t, os.WriteFile(path, []byte(data), 0o600))
  34         }
  35
  36         memory, whatsapp, subagents := backupStartDaemon(t, dir)
  37         require.True(t, whatsapp.isLinked())
  38
  39         _, err := memory.Append("note", "A durable decision.")
  40         require.NoError(t, err)
  41
  42         _, err = memory.Append("note", "The next decision.")
  43         require.NoError(t, err)
  44         require.NoError(t, memory.Settle(t.Context()))
  45
  46         child, err := subagents.Create(t.Context())
  47         require.NoError(t, err)
  48
  49         _, err = child.Append("note", "subagent")
  50         require.NoError(t, err)
  51
  52         _, err = child.Append("user", "A saved task.")
  53         require.NoError(t, err)
  54
  55         childText, err := child.Show()
  56         require.NoError(t, err)
  57
  58         for _, excluded := range []string{"probe/skip", "agent-backup-example/skip"} {
  59                 path := filepath.Join(dir, excluded)
  60                 require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
  61                 require.NoError(t, os.WriteFile(path, []byte("excluded"), 0o600))
  62         }
  63
  64         archive := filepath.Join(t.TempDir(), "daily.tar.gz")
  65         t.Setenv("ONECHAT_DIR", dir)
  66
  67         var answer strings.Builder
  68
  69         require.NoError(t, backup(t.Context(), archive, &answer))
  70         require.Regexp(t, `^`+archive+` \([0-9]+ bytes\)\n$`, answer.String())
  71
  72         info, err := os.Stat(archive)
  73         require.NoError(t, err)
  74         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
  75
  76         restored := t.TempDir()
  77         require.NoError(t, backupRestore(t.Context(), archive, restored))
  78
  79         for name, data := range files {
  80                 bytes, err := os.ReadFile(filepath.Join(restored, name))
  81                 require.NoError(t, err)
  82                 require.Equal(t, data, string(bytes))
  83         }
  84
  85         for _, excluded := range []string{"probe", "agent-backup-example", socketFileName, "agent/lock"} {
  86                 _, err = os.Stat(filepath.Join(restored, excluded))
  87                 require.ErrorIs(t, err, os.ErrNotExist)
  88         }
  89
  90         restoredMemory, restoredWhatsApp, restoredSubagents := backupStartDaemon(t, restored)
  91         require.Equal(t, memory.Len(), restoredMemory.Len())
  92         require.Equal(t, memory.View(), restoredMemory.View())
  93
  94         for id := range memory.Len() {
  95                 expected, err := memory.Message(id)
  96                 require.NoError(t, err)
  97
  98                 actual, err := restoredMemory.Message(id)
  99                 require.NoError(t, err)
 100                 require.Equal(t, expected, actual)
 101         }
 102
 103         require.True(t, restoredWhatsApp.isLinked())
 104         require.Equal(t, whatsapp.link().own, restoredWhatsApp.link().own)
 105
 106         restoredChild, err := restoredSubagents.Get(t.Context(), child.name)
 107         require.NoError(t, err)
 108
 109         restoredText, err := restoredChild.Show()
 110         require.NoError(t, err)
 111         require.Equal(t, childText, restoredText)
 112 }
 113
 114 func TestBackupRestoresEachChildWithItsOwnState(t *testing.T) {
 115         dir := t.TempDir()
 116         daemonChildConfig(t, dir, "Garden")
 117         daemonChildConfig(t, dir, "Library")
 118
 119         server := newModelFake(t)
 120         clock := newTestClock()
 121         agents, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
 122         require.NoError(t, err)
 123
 124         defer func() {
 125                 require.NoError(t, agents.Close())
 126         }()
 127
 128         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
 129
 130         var loops []*TurnLoop
 131
 132         for _, agent := range agents {
 133                 loop, err := agent.Loop(whatsapp, func(Event) {
 134                 })
 135                 require.NoError(t, err)
 136
 137                 loops = append(loops, loop)
 138                 _, err = loop.Submit("parent", "Keep the decision of "+agent.agent.config.Name)
 139                 require.NoError(t, err)
 140                 require.NoError(t, loop.turn(t.Context()))
 141                 require.NoError(t, agent.memory.Settle(t.Context()))
 142
 143                 subagent, err := agent.subagents.Create(t.Context())
 144                 require.NoError(t, err)
 145
 146                 _, err = subagent.Append("user", "Keep the task of "+agent.agent.config.Name)
 147                 require.NoError(t, err)
 148         }
 149
 150         archive := filepath.Join(t.TempDir(), "children.tar.gz")
 151         _, err = backupData(t.Context(), dir, archive, whatsapp, func(ctx context.Context, snapshot func() error) error {
 152                 return daemonSnapshot(ctx, loops, snapshot)
 153         })
 154         require.NoError(t, err)
 155
 156         restored := t.TempDir()
 157         require.NoError(t, backupRestore(t.Context(), archive, restored))
 158
 159         restoredAgents, err := openDaemonAgents(t.Context(), restored, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
 160         require.NoError(t, err)
 161
 162         defer func() {
 163                 require.NoError(t, restoredAgents.Close())
 164         }()
 165
 166         require.Len(t, restoredAgents, 3)
 167
 168         for i, agent := range agents {
 169                 other := restoredAgents[i]
 170                 require.Equal(t, agent.agent.config, other.agent.config)
 171                 require.Equal(t, agent.memory.View(), other.memory.View())
 172                 require.Equal(t, agent.memory.Len(), other.memory.Len())
 173                 require.Equal(t, modelReadUsage(t, agent.client.UsagePath()), modelReadUsage(t, other.client.UsagePath()))
 174
 175                 for id := range agent.memory.Len() {
 176                         before, err := agent.memory.Message(id)
 177                         require.NoError(t, err)
 178
 179                         after, err := other.memory.Message(id)
 180                         require.NoError(t, err)
 181                         require.Equal(t, before, after)
 182                 }
 183
 184                 entries, err := os.ReadDir(filepath.Join(agent.agent.dir, subagentDirName))
 185                 require.NoError(t, err)
 186                 require.Len(t, entries, 1)
 187
 188                 subagent, err := other.subagents.Get(t.Context(), entries[0].Name())
 189                 require.NoError(t, err)
 190
 191                 text, err := subagent.Show()
 192                 require.NoError(t, err)
 193                 require.Contains(t, text, "Keep the task of "+agent.agent.config.Name)
 194         }
 195 }
 196
 197 func TestRestoreRefusesNonemptyDataWithoutChangingIt(t *testing.T) {
 198         dir := t.TempDir()
 199         path := filepath.Join(dir, "keep")
 200         require.NoError(t, os.WriteFile(path, []byte("keep"), 0o600))
 201
 202         err := backupRestore(t.Context(), "missing.tar.gz", dir)
 203         require.ErrorContains(t, err, "empty data directory")
 204
 205         entries, err := os.ReadDir(dir)
 206         require.NoError(t, err)
 207         require.Len(t, entries, 1)
 208
 209         data, err := os.ReadFile(path)
 210         require.NoError(t, err)
 211         require.Equal(t, "keep", string(data))
 212 }
 213
 214 func TestRestoreRefusesTheDaemonTokenLock(t *testing.T) {
 215         dir := t.TempDir()
 216         tokens, err := lockTokens(dir)
 217         require.NoError(t, err)
 218
 219         defer func() {
 220                 require.NoError(t, tokens.Close())
 221         }()
 222
 223         err = backupRestore(t.Context(), "missing.tar.gz", dir)
 224         require.ErrorIs(t, err, errTokensInUse)
 225 }
 226
 227 func TestUnsafeRestoreInstallsNoArchiveFiles(t *testing.T) {
 228         archive := backupTestArchive(t,
 229                 tar.Header{
 230                         Name:     "safe",
 231                         Typeflag: tar.TypeReg,
 232                         Mode:     0o600,
 233                 },
 234                 tar.Header{
 235                         Name:     "../escape",
 236                         Typeflag: tar.TypeReg,
 237                         Mode:     0o600,
 238                 },
 239         )
 240
 241         dir := t.TempDir()
 242         err := backupRestore(t.Context(), archive, dir)
 243         require.ErrorContains(t, err, "unsafe archive path")
 244
 245         entries, err := os.ReadDir(dir)
 246         require.NoError(t, err)
 247         require.Len(t, entries, 1)
 248         require.Equal(t, tokenLockFileName, entries[0].Name())
 249 }
 250
 251 func TestBackupNeedsARunningDaemon(t *testing.T) {
 252         t.Setenv("ONECHAT_DIR", t.TempDir())
 253
 254         err := backup(t.Context(), filepath.Join(t.TempDir(), "backup.tar.gz"), io.Discard)
 255         require.EqualError(t, err, "backup needs a running daemon")
 256 }
 257
 258 func TestWhatsAppSessionBackupIncludesCommittedWALWrites(t *testing.T) {
 259         dir := t.TempDir()
 260         whatsappTestSession(t, dir, true)
 261
 262         whatsapp, _, _ := whatsappStartTest(t, dir, newTestClock())
 263         address, err := whatsappSessionAddress(filepath.Join(dir, whatsappDirName, whatsappSessionFileName))
 264         require.NoError(t, err)
 265
 266         database, err := sql.Open("sqlite", address)
 267         require.NoError(t, err)
 268
 269         defer func() {
 270                 require.NoError(t, database.Close())
 271         }()
 272
 273         _, err = database.ExecContext(t.Context(), "PRAGMA journal_mode=WAL; CREATE TABLE backup_sample (value TEXT); INSERT INTO backup_sample VALUES ('saved');")
 274         require.NoError(t, err)
 275
 276         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 277         require.NoError(t, whatsapp.BackupSession(t.Context(), target))
 278
 279         address, err = whatsappSessionAddress(target)
 280         require.NoError(t, err)
 281
 282         restored, err := sql.Open("sqlite", address)
 283         require.NoError(t, err)
 284
 285         defer func() {
 286                 require.NoError(t, restored.Close())
 287         }()
 288
 289         var value string
 290
 291         err = restored.QueryRowContext(t.Context(), "SELECT value FROM backup_sample").Scan(&value)
 292         require.NoError(t, err)
 293         require.Equal(t, "saved", value)
 294
 295         info, err := os.Stat(target)
 296         require.NoError(t, err)
 297         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 298 }
 299
 300 func TestSessionBackupWaitsForRemovalToFinish(t *testing.T) {
 301         dir := t.TempDir()
 302         whatsappTestSession(t, dir, true)
 303
 304         whatsapp, fake, _ := whatsappStartTest(t, dir, newTestClock())
 305         disconnect := &BackupDisconnect{
 306                 WhatsAppClient: fake,
 307                 entered:        make(chan struct{}),
 308                 release:        make(chan struct{}),
 309         }
 310
 311         var release sync.Once
 312
 313         defer release.Do(func() {
 314                 close(disconnect.release)
 315         })
 316
 317         whatsapp.mutex.Lock()
 318
 319         whatsapp.client = disconnect
 320         whatsapp.mutex.Unlock()
 321
 322         removed := make(chan error, 1)
 323         go func() {
 324                 removed <- whatsapp.removeSession()
 325         }()
 326
 327         <-disconnect.entered
 328
 329         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 330         backedUp := make(chan error, 1)
 331         started := make(chan struct{})
 332         go func() {
 333                 close(started)
 334
 335                 backedUp <- whatsapp.BackupSession(t.Context(), target)
 336         }()
 337
 338         <-started
 339         require.Never(t, func() bool {
 340                 return len(backedUp) != 0
 341         }, 50*time.Millisecond, time.Millisecond)
 342
 343         release.Do(func() {
 344                 close(disconnect.release)
 345         })
 346         require.NoError(t, <-removed)
 347         require.NoError(t, <-backedUp)
 348
 349         for _, path := range []string{whatsapp.sessionPath(), target} {
 350                 _, err := os.Stat(path)
 351                 require.ErrorIs(t, err, os.ErrNotExist)
 352         }
 353 }
 354
 355 // BackupDisconnect holds session removal in its unlocked client cleanup.
 356 type BackupDisconnect struct {
 357         WhatsAppClient
 358         entered chan struct{}
 359         release chan struct{}
 360 }
 361
 362 func (d *BackupDisconnect) Disconnect() {
 363         close(d.entered)
 364         <-d.release
 365 }
 366
 367 // backupStartDaemon serves fake model and WhatsApp clients on the socket of
 368 // dir. It keeps all state files under dir and holds the daemon token lock.
 369 func backupStartDaemon(t *testing.T, dir string) (*Memory, *WhatsApp, *Subagents) {
 370         t.Helper()
 371
 372         tokens, err := lockTokens(dir)
 373         require.NoError(t, err)
 374
 375         clock := newTestClock()
 376         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
 377         require.NoError(t, err)
 378
 379         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer(), discardLog())
 380         require.NoError(t, err)
 381
 382         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
 383         subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), clock.Now, subagentNameRandom)
 384         schedules, err := openSchedules(filepath.Join(agent.dir, scheduleFileName), "OneChat", clock.Now, discardLog())
 385         require.NoError(t, err)
 386
 387         server := newModelFake(t)
 388         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 389         listener, err := daemonListenSocket(filepath.Join(dir, socketFileName))
 390         require.NoError(t, err)
 391
 392         ctx, cancel := context.WithCancel(t.Context())
 393         done := make(chan error, 1)
 394         go func() {
 395                 done <- daemonServe(ctx, listener, nil, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
 396                         return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, clock.Now, discardLog(), notify)
 397                 })
 398         }()
 399
 400         t.Cleanup(func() {
 401                 cancel()
 402                 require.NoError(t, <-done)
 403                 require.NoError(t, subagents.Close())
 404                 require.NoError(t, memory.Close())
 405                 require.NoError(t, tokens.Close())
 406         })
 407
 408         return memory, whatsapp, subagents
 409 }