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 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:  "user message not shown",
 158                         event: Event{Kind: "user", Text: "typed"},
 159                         text:  "",
 160                 },
 161                 {
 162                         name:  "long result cut",
 163                         event: Event{Kind: "echo", Text: strings.Repeat("line\n", 12)},
 164                         text:  "[result] " + strings.Repeat("line\n", 9) + "line\n[2 more lines]",
 165                 },
 166         }
 167
 168         for _, tc := range testCases {
 169                 t.Run(tc.name, func(t *testing.T) {
 170                         require.Equal(t, tc.text, eventText(tc.event))
 171                 })
 172         }
 173 }
 174
 175 // testClient is a terminal client of a test daemon, with a pipe as its input.
 176 type testClient struct {
 177         input *io.PipeWriter
 178         done  chan error
 179         out   *syncBuffer
 180 }
 181
 182 // Type writes one line to the input of the client.
 183 func (c *testClient) Type(line string) {
 184         _, err := io.WriteString(c.input, line+"\n")
 185         if err != nil {
 186                 panic(err)
 187         }
 188 }
 189
 190 // Output returns what the client showed.
 191 func (c *testClient) Output() string {
 192         return c.out.String()
 193 }
 194
 195 // WaitFor waits until the client showed text.
 196 func (c *testClient) WaitFor(t *testing.T, text string) {
 197         t.Helper()
 198
 199         require.Eventually(t, func() bool {
 200                 return strings.Contains(c.Output(), text)
 201         }, 10*time.Second, time.Millisecond, "output so far: %q", c.Output())
 202 }
 203
 204 // Close ends the input and waits for the client.
 205 func (c *testClient) Close() error {
 206         err := c.input.Close()
 207         if err != nil {
 208                 return err
 209         }
 210
 211         return <-c.done
 212 }
 213
 214 // syncBuffer is a buffer that the test and the client use at once.
 215 type syncBuffer struct {
 216         mutex  sync.Mutex
 217         buffer bytes.Buffer
 218 }
 219
 220 func (b *syncBuffer) Write(p []byte) (int, error) {
 221         b.mutex.Lock()
 222         defer b.mutex.Unlock()
 223
 224         return b.buffer.Write(p)
 225 }
 226
 227 func (b *syncBuffer) String() string {
 228         b.mutex.Lock()
 229         defer b.mutex.Unlock()
 230
 231         return b.buffer.String()
 232 }
 233
 234 // startTestDaemon serves a new agent on a socket in a temporary directory and
 235 // returns the socket path. The commands of the agent run in a temporary
 236 // directory.
 237 func startTestDaemon(t *testing.T, server *fakeResponses) string {
 238         t.Helper()
 239
 240         dir := t.TempDir()
 241         agent, err := openAgent(filepath.Join(dir, agentDirName))
 242         require.NoError(t, err)
 243
 244         agent.config.Workdir = t.TempDir()
 245         clock := newTestClock()
 246         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
 247         require.NoError(t, err)
 248
 249         socket := filepath.Join(dir, socketFileName)
 250         listener, err := listenSocket(socket)
 251         require.NoError(t, err)
 252
 253         client := newClient(server.URL, apiKey(testKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 254         ctx, cancel := context.WithCancel(t.Context())
 255         done := make(chan error)
 256         go func() {
 257                 done <- serveAgent(ctx, listener, func(notify func(Event)) (*loop, error) {
 258                         return newLoop(agent, memory, client, clock.Now, notify)
 259                 })
 260         }()
 261
 262         t.Cleanup(func() {
 263                 cancel()
 264                 require.NoError(t, <-done)
 265                 require.NoError(t, memory.Close())
 266         })
 267
 268         return socket
 269 }
 270
 271 // startTestClient connects a client to the daemon at socket.
 272 func startTestClient(t *testing.T, socket string) *testClient {
 273         t.Helper()
 274
 275         conn, err := net.Dial("unix", socket)
 276         require.NoError(t, err)
 277
 278         reader, writer := io.Pipe()
 279         client := &testClient{
 280                 input: writer,
 281                 done:  make(chan error, 1),
 282                 out:   &syncBuffer{},
 283         }
 284
 285         go func() {
 286                 client.done <- chatSession(t.Context(), conn, reader, client.out)
 287         }()
 288
 289         t.Cleanup(func() {
 290                 _ = writer.Close()
 291         })
 292
 293         return client
 294 }