git.giacomolli.uk - onechat/blob - daemon_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
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 func TestClientsSeeTheTurnAsItHappens(t *testing.T) {
  16         server := newModelFake(t)
  17         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  18                 if call.number == 0 {
  19                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "shell", `{"command":"echo hi"}`)}, 100, 0, 5))
  20                         return
  21                 }
  22
  23                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("It said hi.")}, 100, 0, 5))
  24         }
  25
  26         socket := startTestDaemon(t, server)
  27         first := terminalStartTestClient(t, socket)
  28         second := terminalStartTestClient(t, socket)
  29         first.WaitFor(t, "[connected to OneChat]")
  30         second.WaitFor(t, "[connected to OneChat]")
  31
  32         first.Type("Run echo hi.")
  33         for _, client := range []*TerminalTestClient{first, second} {
  34                 client.WaitFor(t, "[done]")
  35                 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())
  36         }
  37
  38         first.Type(stopCommand)
  39         first.WaitFor(t, "[error] no turn runs")
  40         require.NotContains(t, second.Output(), "no turn runs")
  41         require.NoError(t, first.Close())
  42 }
  43
  44 func TestStopInterruptsTheTurn(t *testing.T) {
  45         server := newModelFake(t)
  46         arrived := make(chan struct{})
  47         release := make(chan struct{})
  48         t.Cleanup(func() {
  49                 close(release)
  50         })
  51
  52         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  53                 close(arrived)
  54                 <-release
  55         }
  56
  57         socket := startTestDaemon(t, server)
  58         client := terminalStartTestClient(t, socket)
  59         client.Type("a long task")
  60         <-arrived
  61
  62         late := terminalStartTestClient(t, socket)
  63         late.WaitFor(t, "[connected to OneChat]\n[working]")
  64
  65         client.Type(stopCommand)
  66         client.WaitFor(t, "[error] interrupted\n[done]")
  67 }
  68
  69 func TestListenSocketReplacesAnOldSocketFile(t *testing.T) {
  70         path := filepath.Join(t.TempDir(), socketFileName)
  71         require.NoError(t, os.WriteFile(path, nil, 0o600))
  72
  73         listener, err := listenSocket(path)
  74         require.NoError(t, err)
  75
  76         info, err := os.Stat(path)
  77         require.NoError(t, err)
  78         require.Equal(t, os.ModeSocket, info.Mode().Type())
  79         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
  80         require.NoError(t, listener.Close())
  81 }
  82
  83 func TestListenSocketRefusesALongPath(t *testing.T) {
  84         _, err := listenSocket(filepath.Join(t.TempDir(), strings.Repeat("d", 100), socketFileName))
  85         require.ErrorContains(t, err, "set ONECHAT_DIR to a shorter path")
  86 }
  87
  88 func TestDaemonWritesNoFileForALongSocketPath(t *testing.T) {
  89         dir := filepath.Join(t.TempDir(), strings.Repeat("d", 100))
  90         t.Setenv("ONECHAT_DIR", dir)
  91
  92         err := daemon(t.Context(), io.Discard)
  93         require.ErrorContains(t, err, "set ONECHAT_DIR to a shorter path")
  94
  95         entries, err := os.ReadDir(dir)
  96         require.NoError(t, err)
  97         require.Empty(t, entries)
  98 }
  99
 100 func TestOneProcessHoldsTheTokens(t *testing.T) {
 101         dir := t.TempDir()
 102         lock, err := lockTokens(dir)
 103         require.NoError(t, err)
 104
 105         _, err = lockTokens(dir)
 106         require.ErrorIs(t, err, errTokensInUse)
 107         require.NoError(t, lock.Close())
 108
 109         lock, err = lockTokens(dir)
 110         require.NoError(t, err)
 111         require.NoError(t, lock.Close())
 112 }
 113
 114 // startTestDaemon serves a new agent on a socket in a temporary directory and
 115 // returns the socket path. The commands of the agent run in a temporary
 116 // directory.
 117 func startTestDaemon(t *testing.T, server *ModelFake) string {
 118         t.Helper()
 119
 120         return startTestDaemonWithBot(t, server, nil)
 121 }
 122
 123 // startTestDaemonWithBot is startTestDaemon with bot, which can be nil.
 124 func startTestDaemonWithBot(t *testing.T, server *ModelFake, bot *TelegramBot) string {
 125         t.Helper()
 126
 127         dir := t.TempDir()
 128         agent, err := openAgent(filepath.Join(dir, agentDirName))
 129         require.NoError(t, err)
 130
 131         agent.config.Workdir = t.TempDir()
 132         clock := newTestClock()
 133         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
 134         require.NoError(t, err)
 135
 136         socket := filepath.Join(dir, socketFileName)
 137         listener, err := listenSocket(socket)
 138         require.NoError(t, err)
 139
 140         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 141         ctx, cancel := context.WithCancel(t.Context())
 142         done := make(chan error)
 143         go func() {
 144                 done <- serveAgent(ctx, listener, bot, func(notify func(Event)) (*loop, error) {
 145                         return newLoop(agent, memory, client, clock.Now, notify)
 146                 })
 147         }()
 148
 149         t.Cleanup(func() {
 150                 cancel()
 151                 require.NoError(t, <-done)
 152                 require.NoError(t, memory.Close())
 153         })
 154
 155         return socket
 156 }