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

   1 package main
   2
   3 import (
   4         "bytes"
   5         "context"
   6         "io"
   7         "net"
   8         "net/http"
   9         "os"
  10         "path/filepath"
  11         "strings"
  12         "sync"
  13         "testing"
  14         "time"
  15
  16         "github.com/stretchr/testify/require"
  17 )
  18
  19 func TestClientsSeeTheTurnAsItHappens(t *testing.T) {
  20         server := newModelFake(t)
  21         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  22                 if call.number == 0 {
  23                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "shell", `{"command":"echo hi"}`)}, 100, 0, 5))
  24                         return
  25                 }
  26
  27                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("It said hi.")}, 100, 0, 5))
  28         }
  29
  30         socket := startTestDaemon(t, server)
  31         first := startTestClient(t, socket)
  32         second := startTestClient(t, socket)
  33         first.WaitFor(t, "[connected to OneChat]")
  34         second.WaitFor(t, "[connected to OneChat]")
  35
  36         first.Type("Run echo hi.")
  37         for _, client := range []*testClient{first, second} {
  38                 client.WaitFor(t, "[done]")
  39                 require.Equal(t, "[connected to OneChat]\n[working]\n[tool] shell {\"command\":\"echo hi\"}\n[result] hi\nIt said hi.\n[done]\n", client.Output())
  40         }
  41
  42         first.Type(stopCommand)
  43         first.WaitFor(t, "[error] no turn runs")
  44         require.NotContains(t, second.Output(), "no turn runs")
  45         require.NoError(t, first.Close())
  46 }
  47
  48 func TestStopInterruptsTheTurn(t *testing.T) {
  49         server := newModelFake(t)
  50         arrived := make(chan struct{})
  51         release := make(chan struct{})
  52         t.Cleanup(func() {
  53                 close(release)
  54         })
  55
  56         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  57                 close(arrived)
  58                 <-release
  59         }
  60
  61         socket := startTestDaemon(t, server)
  62         client := startTestClient(t, socket)
  63         client.Type("a long task")
  64         <-arrived
  65
  66         late := startTestClient(t, socket)
  67         late.WaitFor(t, "[connected to OneChat]\n[working]")
  68
  69         client.Type(stopCommand)
  70         client.WaitFor(t, "[error] interrupted\n[done]")
  71 }
  72
  73 func TestClientEndsWhenTheDaemonEnds(t *testing.T) {
  74         dir := t.TempDir()
  75         listener, err := listenSocket(filepath.Join(dir, socketFileName))
  76         require.NoError(t, err)
  77
  78         accepted := make(chan net.Conn)
  79         go func() {
  80                 conn, err := listener.Accept()
  81                 if err == nil {
  82                         accepted <- conn
  83                 }
  84         }()
  85
  86         conn, err := net.Dial("unix", filepath.Join(dir, socketFileName))
  87         require.NoError(t, err)
  88
  89         in, _ := io.Pipe()
  90         done := make(chan error)
  91         go func() {
  92                 done <- chatSession(t.Context(), conn, in, io.Discard)
  93         }()
  94
  95         require.NoError(t, (<-accepted).Close())
  96         require.ErrorIs(t, <-done, errDaemonClosed)
  97         require.NoError(t, listener.Close())
  98 }
  99
 100 func TestListenSocketReplacesAnOldSocketFile(t *testing.T) {
 101         path := filepath.Join(t.TempDir(), socketFileName)
 102         require.NoError(t, os.WriteFile(path, nil, 0o600))
 103
 104         listener, err := listenSocket(path)
 105         require.NoError(t, err)
 106
 107         info, err := os.Stat(path)
 108         require.NoError(t, err)
 109         require.Equal(t, os.ModeSocket, info.Mode().Type())
 110         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 111         require.NoError(t, listener.Close())
 112 }
 113
 114 func TestListenSocketRefusesALongPath(t *testing.T) {
 115         _, err := listenSocket(filepath.Join(t.TempDir(), strings.Repeat("d", 100), socketFileName))
 116         require.ErrorContains(t, err, "set ONECHAT_DIR to a shorter path")
 117 }
 118
 119 func TestDaemonWritesNoFileForALongSocketPath(t *testing.T) {
 120         dir := filepath.Join(t.TempDir(), strings.Repeat("d", 100))
 121         t.Setenv("ONECHAT_DIR", dir)
 122
 123         err := daemon(t.Context(), io.Discard)
 124         require.ErrorContains(t, err, "set ONECHAT_DIR to a shorter path")
 125
 126         entries, err := os.ReadDir(dir)
 127         require.NoError(t, err)
 128         require.Empty(t, entries)
 129 }
 130
 131 func TestOneProcessHoldsTheTokens(t *testing.T) {
 132         dir := t.TempDir()
 133         lock, err := lockTokens(dir)
 134         require.NoError(t, err)
 135
 136         _, err = lockTokens(dir)
 137         require.ErrorIs(t, err, errTokensInUse)
 138         require.NoError(t, lock.Close())
 139
 140         lock, err = lockTokens(dir)
 141         require.NoError(t, err)
 142         require.NoError(t, lock.Close())
 143 }
 144
 145 func TestEventText(t *testing.T) {
 146         testCases := []struct {
 147                 name  string
 148                 event Event
 149                 text  string
 150         }{
 151                 {
 152                         name:  "reply as it is",
 153                         event: Event{Kind: "reply", Text: "Done."},
 154                         text:  "Done.",
 155                 },
 156                 {
 157                         name:  "own user message not shown",
 158                         event: Event{Kind: "user", Text: "typed", Source: terminalSource},
 159                         text:  "",
 160                 },
 161                 {
 162                         name:  "user message of another source marked",
 163                         event: Event{Kind: "user", Text: "from the phone", Source: "telegram"},
 164                         text:  "[telegram] from the phone",
 165                 },
 166                 {
 167                         name:  "long result cut",
 168                         event: Event{Kind: "echo", Text: strings.Repeat("line\n", 12)},
 169                         text:  "[result] " + strings.Repeat("line\n", 9) + "line\n[2 more lines]",
 170                 },
 171         }
 172
 173         for _, tc := range testCases {
 174                 t.Run(tc.name, func(t *testing.T) {
 175                         require.Equal(t, tc.text, eventText(tc.event))
 176                 })
 177         }
 178 }
 179
 180 // testClient is a terminal client of a test daemon, with a pipe as its input.
 181 type testClient struct {
 182         input *io.PipeWriter
 183         done  chan error
 184         out   *syncBuffer
 185 }
 186
 187 // Type writes one line to the input of the client.
 188 func (c *testClient) Type(line string) {
 189         _, err := io.WriteString(c.input, line+"\n")
 190         if err != nil {
 191                 panic(err)
 192         }
 193 }
 194
 195 // Output returns what the client showed.
 196 func (c *testClient) Output() string {
 197         return c.out.String()
 198 }
 199
 200 // WaitFor waits until the client showed text.
 201 func (c *testClient) WaitFor(t *testing.T, text string) {
 202         t.Helper()
 203
 204         require.Eventually(t, func() bool {
 205                 return strings.Contains(c.Output(), text)
 206         }, 10*time.Second, time.Millisecond, "output so far: %q", c.Output())
 207 }
 208
 209 // Close ends the input and waits for the client.
 210 func (c *testClient) Close() error {
 211         err := c.input.Close()
 212         if err != nil {
 213                 return err
 214         }
 215
 216         return <-c.done
 217 }
 218
 219 // syncBuffer is a buffer that the test and the client use at once.
 220 type syncBuffer struct {
 221         mutex  sync.Mutex
 222         buffer bytes.Buffer
 223 }
 224
 225 func (b *syncBuffer) Write(p []byte) (int, error) {
 226         b.mutex.Lock()
 227         defer b.mutex.Unlock()
 228
 229         return b.buffer.Write(p)
 230 }
 231
 232 func (b *syncBuffer) String() string {
 233         b.mutex.Lock()
 234         defer b.mutex.Unlock()
 235
 236         return b.buffer.String()
 237 }
 238
 239 // startTestDaemon serves a new agent on a socket in a temporary directory and
 240 // returns the socket path. The commands of the agent run in a temporary
 241 // directory.
 242 func startTestDaemon(t *testing.T, server *ModelFake) string {
 243         t.Helper()
 244
 245         return startTestDaemonWithBot(t, server, nil)
 246 }
 247
 248 // startTestDaemonWithBot is startTestDaemon with bot, which can be nil.
 249 func startTestDaemonWithBot(t *testing.T, server *ModelFake, bot *TelegramBot) string {
 250         t.Helper()
 251
 252         dir := t.TempDir()
 253         agent, err := openAgent(filepath.Join(dir, agentDirName))
 254         require.NoError(t, err)
 255
 256         agent.config.Workdir = t.TempDir()
 257         clock := newTestClock()
 258         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
 259         require.NoError(t, err)
 260
 261         socket := filepath.Join(dir, socketFileName)
 262         listener, err := listenSocket(socket)
 263         require.NoError(t, err)
 264
 265         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 266         ctx, cancel := context.WithCancel(t.Context())
 267         done := make(chan error)
 268         go func() {
 269                 done <- serveAgent(ctx, listener, bot, func(notify func(Event)) (*loop, error) {
 270                         return newLoop(agent, memory, client, clock.Now, notify)
 271                 })
 272         }()
 273
 274         t.Cleanup(func() {
 275                 cancel()
 276                 require.NoError(t, <-done)
 277                 require.NoError(t, memory.Close())
 278         })
 279
 280         return socket
 281 }
 282
 283 // startTestClient connects a client to the daemon at socket.
 284 func startTestClient(t *testing.T, socket string) *testClient {
 285         t.Helper()
 286
 287         conn, err := net.Dial("unix", socket)
 288         require.NoError(t, err)
 289
 290         reader, writer := io.Pipe()
 291         client := &testClient{
 292                 input: writer,
 293                 done:  make(chan error, 1),
 294                 out:   &syncBuffer{},
 295         }
 296
 297         go func() {
 298                 client.done <- chatSession(t.Context(), conn, reader, client.out)
 299         }()
 300
 301         t.Cleanup(func() {
 302                 _ = writer.Close()
 303         })
 304
 305         return client
 306 }