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