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

   1 package main
   2
   3 import (
   4         "context"
   5         "io"
   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 TestDaemonOpensAndRunsTheRootAndTwoChildren(t *testing.T) {
  17         dir := t.TempDir()
  18         daemonChildConfig(t, dir, "Garden")
  19         daemonChildConfig(t, dir, "Library")
  20
  21         server := newModelFake(t)
  22         clock := newTestClock()
  23         agents, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
  24         require.NoError(t, err)
  25
  26         t.Cleanup(func() {
  27                 require.NoError(t, agents.Close())
  28         })
  29
  30         require.Len(t, agents, 3)
  31         require.Equal(t, []string{"Garden", "Library"}, agents.ChildNames())
  32
  33         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  34                 text := ""
  35                 if call.body["prompt_cache_key"] == agents[0].agent.CacheKey() {
  36                         text = "The decision is recorded."
  37                 }
  38
  39                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(text)}, 100, 0, 5))
  40         }
  41
  42         for _, agent := range agents {
  43                 at := clock.Now().Add(time.Minute).Format(scheduleAtLayout)
  44                 _, err = agent.schedules.Set("Record the decision for "+agent.agent.config.Name, &at, nil)
  45                 require.NoError(t, err)
  46         }
  47
  48         clock.now = clock.now.Add(time.Minute)
  49         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
  50         listener, err := daemonListenSocket(filepath.Join(dir, socketFileName))
  51         require.NoError(t, err)
  52
  53         ctx, cancel := context.WithCancel(t.Context())
  54         defer cancel()
  55
  56         done := make(chan error, 1)
  57         go func() {
  58                 done <- daemonServe(ctx, listener, nil, whatsapp, agents[0].schedules, func(notify func(Event)) (*TurnLoop, error) {
  59                         return agents[0].Loop(whatsapp, notify)
  60                 }, agents[1:]...)
  61         }()
  62
  63         t.Cleanup(func() {
  64                 cancel()
  65                 require.NoError(t, <-done)
  66         })
  67
  68         require.Eventually(t, func() bool {
  69                 for _, agent := range agents {
  70                         count := 1
  71                         if agent == agents[0] {
  72                                 count = 2
  73                         }
  74
  75                         usage, err := os.ReadFile(filepath.Join(agent.agent.dir, "usage.jsonl"))
  76                         if err != nil || !strings.HasSuffix(string(usage), "\n") || agent.memory.Len() != count {
  77                                 return false
  78                         }
  79                 }
  80
  81                 return true
  82         }, 5*time.Second, time.Millisecond)
  83
  84         require.Equal(t, 3, server.Calls())
  85
  86         keys := make(map[string]bool)
  87         for i := range server.Calls() {
  88                 call := server.Call(i)
  89                 keys[call.body["prompt_cache_key"].(string)] = true
  90         }
  91
  92         require.Len(t, keys, 3)
  93
  94         for _, agent := range agents {
  95                 message, err := agent.memory.Message(0)
  96                 require.NoError(t, err)
  97                 require.Contains(t, message.Text, "Record the decision for "+agent.agent.config.Name)
  98                 require.True(t, keys[agent.agent.CacheKey()])
  99
 100                 usage := modelReadUsage(t, filepath.Join(agent.agent.dir, "usage.jsonl"))
 101                 require.Len(t, usage, 1)
 102                 require.Equal(t, "first", usage[0].Turn)
 103                 require.Equal(t, agent.agent.config.Model, usage[0].Model)
 104         }
 105
 106         terminal := terminalStartTestClient(t, listener.Addr().String())
 107         terminal.WaitFor(t, "[connected to OneChat]")
 108         require.NotContains(t, terminal.Output(), "Garden")
 109         require.NotContains(t, terminal.Output(), "Library")
 110 }
 111
 112 func TestDaemonReadyLineNamesTheChildren(t *testing.T) {
 113         dir := t.TempDir()
 114         daemonChildConfig(t, dir, "Garden")
 115         daemonChildConfig(t, dir, "Library")
 116         t.Setenv("ONECHAT_DIR", dir)
 117
 118         var out SyncBuffer
 119
 120         ctx, cancel := context.WithCancel(t.Context())
 121         defer cancel()
 122
 123         done := make(chan error, 1)
 124         go func() {
 125                 done <- daemon(ctx, &out)
 126         }()
 127
 128         t.Cleanup(func() {
 129                 cancel()
 130
 131                 require.NoError(t, <-done)
 132         })
 133
 134         require.Eventually(t, func() bool {
 135                 return strings.Contains(out.String(), "children=[Garden Library]")
 136         }, 5*time.Second, time.Millisecond)
 137 }
 138
 139 func TestChildCallsShareTheirPromptAndToolsWithCompaction(t *testing.T) {
 140         dir := t.TempDir()
 141         daemonChildConfig(t, dir, "Garden")
 142
 143         server := newModelFake(t)
 144         agents, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, time.Now, discardLog())
 145         require.NoError(t, err)
 146
 147         defer func() {
 148                 require.NoError(t, agents.Close())
 149         }()
 150
 151         child := agents[1]
 152         loop, err := child.Loop(nil, func(Event) {
 153         })
 154         require.NoError(t, err)
 155         require.Nil(t, loop.tools.whatsapp)
 156         require.Contains(t, loop.tools.run(t.Context(), "whatsapp_send", `{"id":"fake","chat":"fake","text":"draft"}`), "unknown tool")
 157
 158         _, err = loop.Submit("parent", strings.Repeat("Keep a garden decision. ", nodeBytes))
 159         require.NoError(t, err)
 160         require.NoError(t, loop.turn(t.Context()))
 161         require.NoError(t, child.memory.Settle(t.Context()))
 162
 163         require.GreaterOrEqual(t, server.Calls(), 2)
 164
 165         first := server.Call(0)
 166
 167         for i := range server.Calls() {
 168                 call := server.Call(i)
 169                 require.Equal(t, child.agent.CacheKey(), call.body["prompt_cache_key"])
 170                 require.Equal(t, first.body["instructions"], call.body["instructions"])
 171                 require.Equal(t, first.body["tools"], call.body["tools"])
 172                 require.NotContains(t, string(call.raw), "whatsapp_send")
 173                 require.NotContains(t, string(call.raw), "whatsapp_find")
 174         }
 175
 176         usage := modelReadUsage(t, child.client.UsagePath())
 177         require.GreaterOrEqual(t, len(usage), 2)
 178
 179         var turns, nodes int
 180
 181         for _, record := range usage {
 182                 if record.Node != nil {
 183                         nodes++
 184                 } else {
 185                         turns++
 186                 }
 187         }
 188
 189         require.Positive(t, nodes)
 190         require.Equal(t, 1, turns)
 191 }
 192
 193 func TestBrokenChildConfigStopsDaemonStartup(t *testing.T) {
 194         dir := t.TempDir()
 195         path := filepath.Join(dir, daemonAgentDirName, "children", "Garden")
 196         require.NoError(t, os.MkdirAll(path, 0o700))
 197         require.NoError(t, os.WriteFile(filepath.Join(path, "config.json"), []byte(`{"broken":true}`), 0o600))
 198         t.Setenv("ONECHAT_DIR", dir)
 199
 200         err := daemon(t.Context(), io.Discard)
 201         require.ErrorContains(t, err, `open child "Garden"`)
 202         require.ErrorContains(t, err, `unknown field "broken"`)
 203
 204         _, err = os.Stat(filepath.Join(dir, socketFileName))
 205         require.ErrorIs(t, err, os.ErrNotExist)
 206 }
 207
 208 func TestInvalidChildRelationshipsStopStartup(t *testing.T) {
 209         server := newModelFake(t)
 210         testCases := []struct {
 211                 name   string
 212                 change func(*testing.T, string)
 213                 cause  string
 214         }{
 215                 {
 216                         name: "missing config is not created",
 217                         change: func(t *testing.T, path string) {
 218                                 require.NoError(t, os.Remove(filepath.Join(path, "config.json")))
 219                         },
 220                         cause: `child "Garden" config:`,
 221                 },
 222                 {
 223                         name: "wrong parent is rejected",
 224                         change: func(t *testing.T, path string) {
 225                                 agent, err := openAgent(path)
 226                                 require.NoError(t, err)
 227
 228                                 config := agent.config
 229                                 config.Parent = "Other"
 230                                 agentWriteTestConfig(t, path, config)
 231                         },
 232                         cause: `needs name "Garden" and parent "OneChat"`,
 233                 },
 234                 {
 235                         name: "a grandchild is rejected",
 236                         change: func(t *testing.T, path string) {
 237                                 require.NoError(t, os.MkdirAll(filepath.Join(path, "children", "Nested"), 0o700))
 238                         },
 239                         cause: "children cannot have children",
 240                 },
 241                 {
 242                         name: "a reused ID is rejected",
 243                         change: func(t *testing.T, path string) {
 244                                 root, err := openAgent(filepath.Dir(filepath.Dir(path)))
 245                                 require.NoError(t, err)
 246
 247                                 child, err := openAgent(path)
 248                                 require.NoError(t, err)
 249
 250                                 config := child.config
 251                                 config.ID = root.config.ID
 252                                 agentWriteTestConfig(t, path, config)
 253                         },
 254                         cause: "must have its own name and ID",
 255                 },
 256         }
 257
 258         for _, tc := range testCases {
 259                 t.Run(tc.name, func(t *testing.T) {
 260                         dir := t.TempDir()
 261                         path := daemonChildConfig(t, dir, "Garden")
 262                         tc.change(t, path)
 263
 264                         _, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, time.Now, discardLog())
 265                         require.ErrorContains(t, err, tc.cause)
 266                 })
 267         }
 268 }