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 TestBackupKeepsPendingNotificationsAndCompletedRetryReceipts(t *testing.T) {
 115         directory := t.TempDir()
 116         server := newModelFake(t)
 117         test := newDaemonMessageTest(t, directory, server)
 118         _, err := test.messages.Notify(&NotifyInput{
 119                 Source: "check",
 120                 ID:     "completed",
 121                 Text:   "A completed result.",
 122         })
 123         require.NoError(t, err)
 124         require.NoError(t, test.messages.Notifications(test.loops[0]))
 125         require.NoError(t, test.loops[0].turn(t.Context()))
 126
 127         _, err = test.messages.Notify(&NotifyInput{
 128                 Source: "check",
 129                 ID:     "pending",
 130                 Text:   "A pending result.",
 131         })
 132         require.NoError(t, err)
 133
 134         clock := newTestClock()
 135         whatsapp, _, _ := whatsappStartTest(t, directory, clock)
 136         archive := filepath.Join(t.TempDir(), "events.tar.gz")
 137         _, err = backupData(t.Context(), directory, archive, whatsapp, test.messages.SnapshotAgents)
 138         require.NoError(t, err)
 139
 140         restored := t.TempDir()
 141         require.NoError(t, backupRestore(t.Context(), archive, restored))
 142
 143         before, err := os.ReadFile(test.messages.path)
 144         require.NoError(t, err)
 145
 146         after, err := os.ReadFile(filepath.Join(restored, daemonAgentDirName, "messages.json"))
 147         require.NoError(t, err)
 148         require.Equal(t, before, after)
 149
 150         recovered := newDaemonMessageTest(t, restored, server)
 151         require.Len(t, recovered.Queue(t).Pending, 1)
 152         require.Len(t, recovered.Queue(t).Notifications, 2)
 153
 154         duplicate, err := recovered.messages.Notify(&NotifyInput{
 155                 Source: "check",
 156                 ID:     "completed",
 157                 Text:   "A completed result.",
 158         })
 159         require.NoError(t, err)
 160         require.Equal(t, "duplicate", duplicate.Status)
 161 }
 162
 163 func TestBackupRestoresEachChildWithItsOwnState(t *testing.T) {
 164         dir := t.TempDir()
 165         daemonChildConfig(t, dir, "Garden")
 166         daemonChildConfig(t, dir, "Library")
 167
 168         server := newModelFake(t)
 169         clock := newTestClock()
 170         agents, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
 171         require.NoError(t, err)
 172
 173         defer func() {
 174                 require.NoError(t, agents.Close())
 175         }()
 176
 177         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
 178
 179         var loops []*TurnLoop
 180
 181         for _, agent := range agents {
 182                 loop, err := agent.Loop(whatsapp, func(Event) {
 183                 })
 184                 require.NoError(t, err)
 185
 186                 loops = append(loops, loop)
 187                 _, err = loop.Submit("parent", "Keep the decision of "+agent.agent.config.Name)
 188                 require.NoError(t, err)
 189                 require.NoError(t, loop.turn(t.Context()))
 190                 require.NoError(t, agent.memory.Settle(t.Context()))
 191
 192                 subagent, err := agent.subagents.Create(t.Context())
 193                 require.NoError(t, err)
 194
 195                 _, err = subagent.Append("user", "Keep the task of "+agent.agent.config.Name)
 196                 require.NoError(t, err)
 197         }
 198
 199         archive := filepath.Join(t.TempDir(), "children.tar.gz")
 200         _, err = backupData(t.Context(), dir, archive, whatsapp, func(ctx context.Context, snapshot func() error) error {
 201                 return daemonSnapshot(ctx, loops, snapshot)
 202         })
 203         require.NoError(t, err)
 204
 205         restored := t.TempDir()
 206         require.NoError(t, backupRestore(t.Context(), archive, restored))
 207
 208         restoredAgents, err := openDaemonAgents(t.Context(), restored, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
 209         require.NoError(t, err)
 210
 211         defer func() {
 212                 require.NoError(t, restoredAgents.Close())
 213         }()
 214
 215         require.Len(t, restoredAgents, 3)
 216
 217         for i, agent := range agents {
 218                 other := restoredAgents[i]
 219                 require.Equal(t, agent.agent.config, other.agent.config)
 220                 require.Equal(t, agent.memory.View(), other.memory.View())
 221                 require.Equal(t, agent.memory.Len(), other.memory.Len())
 222                 require.Equal(t, modelReadUsage(t, agent.client.UsagePath()), modelReadUsage(t, other.client.UsagePath()))
 223
 224                 for id := range agent.memory.Len() {
 225                         before, err := agent.memory.Message(id)
 226                         require.NoError(t, err)
 227
 228                         after, err := other.memory.Message(id)
 229                         require.NoError(t, err)
 230                         require.Equal(t, before, after)
 231                 }
 232
 233                 entries, err := os.ReadDir(filepath.Join(agent.agent.dir, subagentDirName))
 234                 require.NoError(t, err)
 235                 require.Len(t, entries, 1)
 236
 237                 subagent, err := other.subagents.Get(t.Context(), entries[0].Name())
 238                 require.NoError(t, err)
 239
 240                 text, err := subagent.Show()
 241                 require.NoError(t, err)
 242                 require.Contains(t, text, "Keep the task of "+agent.agent.config.Name)
 243         }
 244 }
 245
 246 func TestRestoreRefusesNonemptyDataWithoutChangingIt(t *testing.T) {
 247         dir := t.TempDir()
 248         path := filepath.Join(dir, "keep")
 249         require.NoError(t, os.WriteFile(path, []byte("keep"), 0o600))
 250
 251         err := backupRestore(t.Context(), "missing.tar.gz", dir)
 252         require.ErrorContains(t, err, "empty data directory")
 253
 254         entries, err := os.ReadDir(dir)
 255         require.NoError(t, err)
 256         require.Len(t, entries, 1)
 257
 258         data, err := os.ReadFile(path)
 259         require.NoError(t, err)
 260         require.Equal(t, "keep", string(data))
 261 }
 262
 263 func TestRestoreRefusesTheDaemonTokenLock(t *testing.T) {
 264         dir := t.TempDir()
 265         tokens, err := lockTokens(dir)
 266         require.NoError(t, err)
 267
 268         defer func() {
 269                 require.NoError(t, tokens.Close())
 270         }()
 271
 272         err = backupRestore(t.Context(), "missing.tar.gz", dir)
 273         require.ErrorIs(t, err, errTokensInUse)
 274 }
 275
 276 func TestUnsafeRestoreInstallsNoArchiveFiles(t *testing.T) {
 277         archive := backupTestArchive(t,
 278                 tar.Header{
 279                         Name:     "safe",
 280                         Typeflag: tar.TypeReg,
 281                         Mode:     0o600,
 282                 },
 283                 tar.Header{
 284                         Name:     "../escape",
 285                         Typeflag: tar.TypeReg,
 286                         Mode:     0o600,
 287                 },
 288         )
 289
 290         dir := t.TempDir()
 291         err := backupRestore(t.Context(), archive, dir)
 292         require.ErrorContains(t, err, "unsafe archive path")
 293
 294         entries, err := os.ReadDir(dir)
 295         require.NoError(t, err)
 296         require.Len(t, entries, 1)
 297         require.Equal(t, tokenLockFileName, entries[0].Name())
 298 }
 299
 300 func TestBackupNeedsARunningDaemon(t *testing.T) {
 301         t.Setenv("ONECHAT_DIR", t.TempDir())
 302
 303         err := backup(t.Context(), filepath.Join(t.TempDir(), "backup.tar.gz"), io.Discard)
 304         require.EqualError(t, err, "backup needs a running daemon")
 305 }
 306
 307 func TestWhatsAppSessionBackupIncludesCommittedWALWrites(t *testing.T) {
 308         dir := t.TempDir()
 309         whatsappTestSession(t, dir, true)
 310
 311         whatsapp, _, _ := whatsappStartTest(t, dir, newTestClock())
 312         address, err := whatsappSessionAddress(filepath.Join(dir, whatsappDirName, whatsappSessionFileName))
 313         require.NoError(t, err)
 314
 315         database, err := sql.Open("sqlite", address)
 316         require.NoError(t, err)
 317
 318         defer func() {
 319                 require.NoError(t, database.Close())
 320         }()
 321
 322         _, err = database.ExecContext(t.Context(), "PRAGMA journal_mode=WAL; CREATE TABLE backup_sample (value TEXT); INSERT INTO backup_sample VALUES ('saved');")
 323         require.NoError(t, err)
 324
 325         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 326         require.NoError(t, whatsapp.BackupSession(t.Context(), target))
 327
 328         address, err = whatsappSessionAddress(target)
 329         require.NoError(t, err)
 330
 331         restored, err := sql.Open("sqlite", address)
 332         require.NoError(t, err)
 333
 334         defer func() {
 335                 require.NoError(t, restored.Close())
 336         }()
 337
 338         var value string
 339
 340         err = restored.QueryRowContext(t.Context(), "SELECT value FROM backup_sample").Scan(&value)
 341         require.NoError(t, err)
 342         require.Equal(t, "saved", value)
 343
 344         info, err := os.Stat(target)
 345         require.NoError(t, err)
 346         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 347 }
 348
 349 func TestSessionBackupWaitsForRemovalToFinish(t *testing.T) {
 350         dir := t.TempDir()
 351         whatsappTestSession(t, dir, true)
 352
 353         whatsapp, fake, _ := whatsappStartTest(t, dir, newTestClock())
 354         disconnect := &BackupDisconnect{
 355                 WhatsAppClient: fake,
 356                 entered:        make(chan struct{}),
 357                 release:        make(chan struct{}),
 358         }
 359
 360         var release sync.Once
 361
 362         defer release.Do(func() {
 363                 close(disconnect.release)
 364         })
 365
 366         whatsapp.mutex.Lock()
 367
 368         whatsapp.client = disconnect
 369         whatsapp.mutex.Unlock()
 370
 371         removed := make(chan error, 1)
 372         go func() {
 373                 removed <- whatsapp.removeSession()
 374         }()
 375
 376         <-disconnect.entered
 377
 378         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 379         backedUp := make(chan error, 1)
 380         started := make(chan struct{})
 381         go func() {
 382                 close(started)
 383
 384                 backedUp <- whatsapp.BackupSession(t.Context(), target)
 385         }()
 386
 387         <-started
 388         require.Never(t, func() bool {
 389                 return len(backedUp) != 0
 390         }, 50*time.Millisecond, time.Millisecond)
 391
 392         release.Do(func() {
 393                 close(disconnect.release)
 394         })
 395         require.NoError(t, <-removed)
 396         require.NoError(t, <-backedUp)
 397
 398         for _, path := range []string{whatsapp.sessionPath(), target} {
 399                 _, err := os.Stat(path)
 400                 require.ErrorIs(t, err, os.ErrNotExist)
 401         }
 402 }
 403
 404 // BackupDisconnect holds session removal in its unlocked client cleanup.
 405 type BackupDisconnect struct {
 406         WhatsAppClient
 407         entered chan struct{}
 408         release chan struct{}
 409 }
 410
 411 func (d *BackupDisconnect) Disconnect() {
 412         close(d.entered)
 413         <-d.release
 414 }
 415
 416 // backupStartDaemon serves fake model and WhatsApp clients on the socket of
 417 // dir. It keeps all state files under dir and holds the daemon token lock.
 418 func backupStartDaemon(t *testing.T, dir string) (*Memory, *WhatsApp, *Subagents) {
 419         t.Helper()
 420
 421         tokens, err := lockTokens(dir)
 422         require.NoError(t, err)
 423
 424         clock := newTestClock()
 425         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
 426         require.NoError(t, err)
 427
 428         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer(), discardLog())
 429         require.NoError(t, err)
 430
 431         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
 432         subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), clock.Now, subagentNameRandom)
 433         schedules, err := openSchedules(filepath.Join(agent.dir, scheduleFileName), "OneChat", clock.Now, discardLog())
 434         require.NoError(t, err)
 435
 436         server := newModelFake(t)
 437         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 438         listener, err := daemonListenSocket(filepath.Join(dir, socketFileName))
 439         require.NoError(t, err)
 440
 441         ctx, cancel := context.WithCancel(t.Context())
 442         done := make(chan error, 1)
 443         go func() {
 444                 done <- daemonServe(ctx, listener, nil, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
 445                         return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, clock.Now, discardLog(), notify)
 446                 })
 447         }()
 448
 449         t.Cleanup(func() {
 450                 cancel()
 451                 require.NoError(t, <-done)
 452                 require.NoError(t, subagents.Close())
 453                 require.NoError(t, memory.Close())
 454                 require.NoError(t, tokens.Close())
 455         })
 456
 457         return memory, whatsapp, subagents
 458 }