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