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 := newFakeResponses(t)
  21         server.reply = func(w http.ResponseWriter, call fakeCall) {
  22                 if call.number == 0 {
  23                         writeEvents(w, completedEvent([]any{functionCallItem("fc_1", "call_1", "shell", `{"command":"echo hi"}`)}, 100, 0, 5))
  24                         return
  25                 }
  26
  27                 writeEvents(w, completedEvent([]any{messageItem("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 := newFakeResponses(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 fakeCall) {
  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 TestOneProcessHoldsTheTokens(t *testing.T) {
 120         dir := t.TempDir()
 121         lock, err := lockTokens(dir)
 122         require.NoError(t, err)
 123
 124         _, err = lockTokens(dir)
 125         require.ErrorIs(t, err, errTokensInUse)
 126         require.NoError(t, lock.Close())
 127
 128         lock, err = lockTokens(dir)
 129         require.NoError(t, err)
 130         require.NoError(t, lock.Close())
 131 }
 132
 133 func TestEventText(t *testing.T) {
 134         testCases := []struct {
 135                 name  string
 136                 event Event
 137                 text  string
 138         }{
 139                 {
 140                         name:  "reply as it is",
 141                         event: Event{Kind: "reply", Text: "Done."},
 142                         text:  "Done.",
 143                 },
 144                 {
 145                         name:  "user message not shown",
 146                         event: Event{Kind: "user", Text: "typed"},
 147                         text:  "",
 148                 },
 149                 {
 150                         name:  "long result cut",
 151                         event: Event{Kind: "echo", Text: strings.Repeat("line\n", 12)},
 152                         text:  "[result] " + strings.Repeat("line\n", 9) + "line\n[2 more lines]",
 153                 },
 154         }
 155
 156         for _, tc := range testCases {
 157                 t.Run(tc.name, func(t *testing.T) {
 158                         require.Equal(t, tc.text, eventText(tc.event))
 159                 })
 160         }
 161 }
 162
 163 // testClient is a terminal client of a test daemon, with a pipe as its input.
 164 type testClient struct {
 165         input *io.PipeWriter
 166         done  chan error
 167         out   *syncBuffer
 168 }
 169
 170 // Type writes one line to the input of the client.
 171 func (c *testClient) Type(line string) {
 172         _, err := io.WriteString(c.input, line+"\n")
 173         if err != nil {
 174                 panic(err)
 175         }
 176 }
 177
 178 // Output returns what the client showed.
 179 func (c *testClient) Output() string {
 180         return c.out.String()
 181 }
 182
 183 // WaitFor waits until the client showed text.
 184 func (c *testClient) WaitFor(t *testing.T, text string) {
 185         t.Helper()
 186
 187         require.Eventually(t, func() bool {
 188                 return strings.Contains(c.Output(), text)
 189         }, 10*time.Second, time.Millisecond, "output so far: %q", c.Output())
 190 }
 191
 192 // Close ends the input and waits for the client.
 193 func (c *testClient) Close() error {
 194         err := c.input.Close()
 195         if err != nil {
 196                 return err
 197         }
 198
 199         return <-c.done
 200 }
 201
 202 // syncBuffer is a buffer that the test and the client use at once.
 203 type syncBuffer struct {
 204         mutex  sync.Mutex
 205         buffer bytes.Buffer
 206 }
 207
 208 func (b *syncBuffer) Write(p []byte) (int, error) {
 209         b.mutex.Lock()
 210         defer b.mutex.Unlock()
 211
 212         return b.buffer.Write(p)
 213 }
 214
 215 func (b *syncBuffer) String() string {
 216         b.mutex.Lock()
 217         defer b.mutex.Unlock()
 218
 219         return b.buffer.String()
 220 }
 221
 222 // startTestDaemon serves a new agent on a socket in a temporary directory and
 223 // returns the socket path. The commands of the agent run in a temporary
 224 // directory.
 225 func startTestDaemon(t *testing.T, server *fakeResponses) string {
 226         t.Helper()
 227
 228         dir := t.TempDir()
 229         agent, err := openAgent(filepath.Join(dir, agentDirName))
 230         require.NoError(t, err)
 231
 232         agent.config.Workdir = t.TempDir()
 233         clock := newTestClock()
 234         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
 235         require.NoError(t, err)
 236
 237         socket := filepath.Join(dir, socketFileName)
 238         listener, err := listenSocket(socket)
 239         require.NoError(t, err)
 240
 241         client := newClient(server.URL, apiKey(testKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 242         ctx, cancel := context.WithCancel(t.Context())
 243         done := make(chan error)
 244         go func() {
 245                 done <- serveAgent(ctx, listener, func(notify func(Event)) (*loop, error) {
 246                         return newLoop(agent, memory, client, clock.Now, notify)
 247                 })
 248         }()
 249
 250         t.Cleanup(func() {
 251                 cancel()
 252                 require.NoError(t, <-done)
 253                 require.NoError(t, memory.Close())
 254         })
 255
 256         return socket
 257 }
 258
 259 // startTestClient connects a client to the daemon at socket.
 260 func startTestClient(t *testing.T, socket string) *testClient {
 261         t.Helper()
 262
 263         conn, err := net.Dial("unix", socket)
 264         require.NoError(t, err)
 265
 266         reader, writer := io.Pipe()
 267         client := &testClient{
 268                 input: writer,
 269                 done:  make(chan error, 1),
 270                 out:   &syncBuffer{},
 271         }
 272
 273         go func() {
 274                 client.done <- chatSession(t.Context(), conn, reader, client.out)
 275         }()
 276
 277         t.Cleanup(func() {
 278                 _ = writer.Close()
 279         })
 280
 281         return client
 282 }