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