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

   1 package main
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "os"
   7         "path/filepath"
   8         "strings"
   9         "sync"
  10         "testing"
  11         "time"
  12         "unicode/utf8"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 func TestToolRunsAfterItsCallIsDurable(t *testing.T) {
  18         dir := t.TempDir()
  19         log := &ToolRecordingLog{
  20                 marker: filepath.Join(dir, "ran"),
  21         }
  22
  23         box := newTestToolbox(t, log, dir)
  24         output, err := box.Call(t.Context(), "shell", `{"command":"touch ran && echo done"}`)
  25         require.NoError(t, err)
  26         require.Equal(t, "done\n", output)
  27         require.Equal(t, []string{
  28                 `tool: shell {"command":"touch ran && echo done"}`,
  29                 "flush, marker absent",
  30                 "echo: done\n",
  31         }, log.Events())
  32         require.FileExists(t, filepath.Join(dir, "ran"))
  33 }
  34
  35 func TestToolDoesNotRunAfterAnInterruptBeforeTheFlush(t *testing.T) {
  36         dir := t.TempDir()
  37         log := &ToolRecordingLog{
  38                 marker: filepath.Join(dir, "ran"),
  39         }
  40
  41         ctx, cancel := context.WithCancel(t.Context())
  42         cancel()
  43
  44         box := newTestToolbox(t, log, dir)
  45         output, err := box.Call(ctx, "shell", `{"command":"touch ran"}`)
  46         require.NoError(t, err)
  47         require.Equal(t, "not run: the turn was interrupted", output)
  48         require.NoFileExists(t, filepath.Join(dir, "ran"))
  49         require.Equal(t, "echo: not run: the turn was interrupted", log.Events()[2])
  50 }
  51
  52 func TestShellRunsInTheWorkingDirectory(t *testing.T) {
  53         dir := t.TempDir()
  54         box := newTestToolbox(t, &ToolRecordingLog{}, dir)
  55         testCases := []struct {
  56                 name    string
  57                 command string
  58                 output  string
  59         }{
  60                 {
  61                         name:    "working directory",
  62                         command: "pwd",
  63                         output:  dir + "\n",
  64                 },
  65                 {
  66                         name:    "error output and exit status",
  67                         command: "echo out; echo err >&2; exit 3",
  68                         output:  "out\nerr\n\n[exit status 3]",
  69                 },
  70                 {
  71                         name:    "no output",
  72                         command: "true",
  73                         output:  "(no output)",
  74                 },
  75                 {
  76                         name:    "no input",
  77                         command: "cat",
  78                         output:  "(no output)",
  79                 },
  80         }
  81
  82         for _, tc := range testCases {
  83                 t.Run(tc.name, func(t *testing.T) {
  84                         require.Equal(t, tc.output, box.run(t.Context(), "shell", fmt.Sprintf(`{"command":%q}`, tc.command)))
  85                 })
  86         }
  87 }
  88
  89 func TestShellInterruptKillsTheCommandAndItsChildren(t *testing.T) {
  90         box := newTestToolbox(t, &ToolRecordingLog{}, t.TempDir())
  91         ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond)
  92         defer cancel()
  93
  94         start := time.Now()
  95         output := box.shell(ctx, "echo started; sleep 30 & sleep 30")
  96         require.Less(t, time.Since(start), 10*time.Second)
  97         require.Equal(t, "started\n\n[interrupted]", output)
  98 }
  99
 100 func TestFileTools(t *testing.T) {
 101         dir := t.TempDir()
 102         box := newTestToolbox(t, &ToolRecordingLog{}, dir)
 103         path := filepath.Join(dir, "notes", "a.txt")
 104
 105         require.Equal(t, "wrote 12 bytes to "+path, box.run(t.Context(), "write", `{"path":"notes/a.txt","text":"one two two\n"}`))
 106         require.Equal(t, "one two two\n", box.run(t.Context(), "read", fmt.Sprintf(`{"path":%q}`, path)))
 107         require.Equal(t, "error: old occurs 2 times in "+path+", not once", box.run(t.Context(), "edit", `{"path":"notes/a.txt","old":"two","new":"2"}`))
 108         require.Equal(t, "edited "+path, box.run(t.Context(), "edit", `{"path":"notes/a.txt","old":"one","new":"1"}`))
 109         require.Equal(t, "1 two two\n", box.run(t.Context(), "read", `{"path":"notes/a.txt"}`))
 110         require.Contains(t, box.run(t.Context(), "read", `{"path":"missing.txt"}`), "error: open "+filepath.Join(dir, "missing.txt"))
 111         require.Equal(t, `error: arguments: json: unknown field "file"`, box.run(t.Context(), "read", `{"file":"a.txt"}`))
 112         require.Equal(t, `error: unknown tool "grep"`, box.run(t.Context(), "grep", `{}`))
 113 }
 114
 115 func TestZoomAndDateReadTheChat(t *testing.T) {
 116         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
 117         require.NoError(t, err)
 118
 119         defer func() {
 120                 require.NoError(t, memory.Close())
 121         }()
 122
 123         for _, text := range []string{"zero", "one", "two", "three"} {
 124                 _, err = memory.Append("user", text)
 125                 require.NoError(t, err)
 126         }
 127
 128         box := &Toolbox{
 129                 log:     memory,
 130                 chat:    memory,
 131                 workdir: t.TempDir(),
 132         }
 133
 134         testCases := []struct {
 135                 name      string
 136                 arguments string
 137                 output    string
 138         }{
 139                 {
 140                         name:      "one message whole",
 141                         arguments: `{"id":2,"n":1}`,
 142                         output:    "user: two",
 143                 },
 144                 {
 145                         name:      "a line opens into its two halves",
 146                         arguments: `{"id":0,"n":4}`,
 147                         output:    "0+2|user: zero user: one\n2+2|user: two user: three",
 148                 },
 149                 {
 150                         name:      "an id that is not a multiple of n",
 151                         arguments: `{"id":1,"n":2}`,
 152                         output:    "error: line 1+2 does not exist: n must be a power of 2 and id a multiple of n",
 153                 },
 154                 {
 155                         name:      "a line past the last message",
 156                         arguments: `{"id":4,"n":4}`,
 157                         output:    "error: line 4+4 goes past the last message, 3",
 158                 },
 159         }
 160
 161         for _, tc := range testCases {
 162                 t.Run(tc.name, func(t *testing.T) {
 163                         require.Equal(t, tc.output, box.run(t.Context(), "zoom", tc.arguments))
 164                 })
 165         }
 166
 167         date := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC).Local().Format("Monday 2006-01-02 15:04:05 -07:00")
 168         require.Equal(t, date, box.run(t.Context(), "date", `{"id":3}`))
 169         require.Equal(t, "error: message 4 does not exist", box.run(t.Context(), "date", `{"id":4}`))
 170 }
 171
 172 func TestClipKeepsTheHeadAndTheTail(t *testing.T) {
 173         short := strings.Repeat("a", toolClipCharacters)
 174         require.Equal(t, short, toolClip(short))
 175
 176         text := strings.Repeat("h", toolClipCharacters/2) + strings.Repeat("é", 5000) + strings.Repeat("t", toolClipCharacters/2)
 177         require.Equal(t, strings.Repeat("h", toolClipCharacters/2)+"\n[... 10000 bytes cut ...]\n"+strings.Repeat("t", toolClipCharacters/2), toolClip(text))
 178
 179         half := strings.Repeat("é", toolClipCharacters/2)
 180         clipped := toolClip(strings.Repeat("é", toolClipCharacters+2))
 181         require.True(t, utf8.ValidString(clipped))
 182         require.Equal(t, half+"\n[... 4 bytes cut ...]\n"+half, clipped)
 183 }
 184
 185 func TestHeadTailKeepsTheEndsOfALongOutput(t *testing.T) {
 186         output := &ToolHeadTail{}
 187         for n := range 1000 {
 188                 _, err := fmt.Fprintf(output, "%04d%s\n", n, strings.Repeat("x", 995))
 189                 require.NoError(t, err)
 190         }
 191
 192         text := output.String()
 193         require.True(t, strings.HasPrefix(text, "0000xxx"))
 194         require.True(t, strings.HasSuffix(text, "0999"+strings.Repeat("x", 995)+"\n"))
 195         require.Contains(t, text, fmt.Sprintf("\n[... %d bytes cut ...]\n", 1000*1000-toolClipCharacters))
 196         require.Equal(t, toolClipCharacters+len(fmt.Sprintf("\n[... %d bytes cut ...]\n", 1000*1000-toolClipCharacters)), len(text))
 197 }
 198
 199 // ToolRecordingLog is a chat log that records each append and flush. The flush
 200 // records whether the marker file exists.
 201 type ToolRecordingLog struct {
 202         mutex  sync.Mutex
 203         marker string
 204         events []string
 205 }
 206
 207 func (r *ToolRecordingLog) Append(kind, text string) (Message, error) {
 208         r.mutex.Lock()
 209         defer r.mutex.Unlock()
 210
 211         r.events = append(r.events, kind+": "+text)
 212
 213         return Message{}, nil
 214 }
 215
 216 func (r *ToolRecordingLog) Flush(ctx context.Context) error {
 217         r.mutex.Lock()
 218         defer r.mutex.Unlock()
 219
 220         state := "absent"
 221         _, err := os.Stat(r.marker)
 222         if err == nil {
 223                 state = "present"
 224         }
 225
 226         r.events = append(r.events, "flush, marker "+state)
 227
 228         return ctx.Err()
 229 }
 230
 231 func (r *ToolRecordingLog) Events() []string {
 232         r.mutex.Lock()
 233         defer r.mutex.Unlock()
 234
 235         return append([]string(nil), r.events...)
 236 }
 237
 238 // newTestToolbox returns a toolbox with log, an empty chat, and the working
 239 // directory dir.
 240 func newTestToolbox(t *testing.T, log ToolLog, dir string) *Toolbox {
 241         t.Helper()
 242
 243         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
 244         require.NoError(t, err)
 245
 246         t.Cleanup(func() {
 247                 require.NoError(t, memory.Close())
 248         })
 249
 250         return &Toolbox{
 251                 log:     log,
 252                 chat:    memory,
 253                 workdir: dir,
 254         }
 255 }