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