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 := daemonStartTest(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 := daemonStartTest(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 := daemonListenSocket(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 := daemonListenSocket(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 // daemonStartTest serves a new agent on a socket in a temporary directory and
 101 // returns the socket path. The commands of the agent run in a temporary
 102 // directory.
 103 func daemonStartTest(t *testing.T, server *ModelFake) string {
 104         t.Helper()
 105
 106         return daemonStartTestWithBot(t, server, nil)
 107 }
 108
 109 // daemonStartTestWithBot is daemonStartTest with bot, which can be nil.
 110 func daemonStartTestWithBot(t *testing.T, server *ModelFake, bot *TelegramBot) string {
 111         t.Helper()
 112
 113         dir := t.TempDir()
 114         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
 115         require.NoError(t, err)
 116
 117         agent.config.Workdir = t.TempDir()
 118         clock := newTestClock()
 119         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
 120         require.NoError(t, err)
 121
 122         socket := filepath.Join(dir, socketFileName)
 123         listener, err := daemonListenSocket(socket)
 124         require.NoError(t, err)
 125
 126         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 127         ctx, cancel := context.WithCancel(t.Context())
 128         done := make(chan error)
 129         go func() {
 130                 done <- daemonServe(ctx, listener, bot, func(notify func(Event)) (*TurnLoop, error) {
 131                         return newTurnLoop(agent, memory, client, clock.Now, notify)
 132                 })
 133         }()
 134
 135         t.Cleanup(func() {
 136                 cancel()
 137                 require.NoError(t, <-done)
 138                 require.NoError(t, memory.Close())
 139         })
 140
 141         return socket
 142 }