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                 subagents: newTestSubagents(t),
 132                 workdir:   t.TempDir(),
 133         }
 134
 135         testCases := []struct {
 136                 name      string
 137                 arguments string
 138                 output    string
 139         }{
 140                 {
 141                         name:      "one message whole",
 142                         arguments: `{"id":2,"n":1}`,
 143                         output:    "user: two",
 144                 },
 145                 {
 146                         name:      "a line opens into its two halves",
 147                         arguments: `{"id":0,"n":4}`,
 148                         output:    "0+2|user: zero user: one\n2+2|user: two user: three",
 149                 },
 150                 {
 151                         name:      "an id that is not a multiple of n",
 152                         arguments: `{"id":1,"n":2}`,
 153                         output:    "error: line 1+2 does not exist: n must be a power of 2 and id a multiple of n",
 154                 },
 155                 {
 156                         name:      "a line past the last message",
 157                         arguments: `{"id":4,"n":4}`,
 158                         output:    "error: line 4+4 goes past the last message, 3",
 159                 },
 160         }
 161
 162         for _, tc := range testCases {
 163                 t.Run(tc.name, func(t *testing.T) {
 164                         require.Equal(t, tc.output, box.run(t.Context(), "zoom", tc.arguments))
 165                 })
 166         }
 167
 168         date := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC).Local().Format("Monday 2006-01-02 15:04:05 -07:00")
 169         require.Equal(t, date, box.run(t.Context(), "date", `{"id":3}`))
 170         require.Equal(t, "error: message 4 does not exist", box.run(t.Context(), "date", `{"id":4}`))
 171 }
 172
 173 func TestZoomWithANameGivesTheWholeChatOfTheAgent(t *testing.T) {
 174         box := newTestToolbox(t, &ToolRecordingLog{}, t.TempDir())
 175         subagent, err := box.subagents.Create(t.Context())
 176         require.NoError(t, err)
 177
 178         for _, message := range [][2]string{{"note", "subagent"}, {"user", "Count the files."}} {
 179                 _, err = subagent.Append(message[0], message[1])
 180                 require.NoError(t, err)
 181         }
 182
 183         output := box.run(t.Context(), "zoom", `{"id":0,"n":1,"name":"`+subagent.name+`"}`)
 184         require.Equal(t, "0|note: subagent\n\n1|user: Count the files.", output)
 185
 186         require.Equal(t, "error: no agent BlueFox", box.run(t.Context(), "zoom", `{"id":0,"n":1,"name":"BlueFox"}`))
 187         require.Equal(t, "error: line 0+1 goes past the last message, -1", box.run(t.Context(), "zoom", `{"id":0,"n":1,"name":null}`))
 188 }
 189
 190 func TestClipKeepsTheHeadAndTheTail(t *testing.T) {
 191         short := strings.Repeat("a", toolClipCharacters)
 192         require.Equal(t, short, toolClip(short))
 193
 194         text := strings.Repeat("h", toolClipCharacters/2) + strings.Repeat("é", 5000) + strings.Repeat("t", toolClipCharacters/2)
 195         require.Equal(t, strings.Repeat("h", toolClipCharacters/2)+"\n[... 10000 bytes cut ...]\n"+strings.Repeat("t", toolClipCharacters/2), toolClip(text))
 196
 197         half := strings.Repeat("é", toolClipCharacters/2)
 198         clipped := toolClip(strings.Repeat("é", toolClipCharacters+2))
 199         require.True(t, utf8.ValidString(clipped))
 200         require.Equal(t, half+"\n[... 4 bytes cut ...]\n"+half, clipped)
 201 }
 202
 203 func TestHeadTailKeepsTheEndsOfALongOutput(t *testing.T) {
 204         output := &ToolHeadTail{}
 205         for n := range 1000 {
 206                 _, err := fmt.Fprintf(output, "%04d%s\n", n, strings.Repeat("x", 995))
 207                 require.NoError(t, err)
 208         }
 209
 210         text := output.String()
 211         require.True(t, strings.HasPrefix(text, "0000xxx"))
 212         require.True(t, strings.HasSuffix(text, "0999"+strings.Repeat("x", 995)+"\n"))
 213         require.Contains(t, text, fmt.Sprintf("\n[... %d bytes cut ...]\n", 1000*1000-toolClipCharacters))
 214         require.Equal(t, toolClipCharacters+len(fmt.Sprintf("\n[... %d bytes cut ...]\n", 1000*1000-toolClipCharacters)), len(text))
 215 }
 216
 217 // ToolRecordingLog is a chat log that records each append and flush. The flush
 218 // records whether the marker file exists.
 219 type ToolRecordingLog struct {
 220         mutex  sync.Mutex
 221         marker string
 222         events []string
 223 }
 224
 225 func (r *ToolRecordingLog) Append(kind, text string) (Message, error) {
 226         r.mutex.Lock()
 227         defer r.mutex.Unlock()
 228
 229         r.events = append(r.events, kind+": "+text)
 230
 231         return Message{}, nil
 232 }
 233
 234 func (r *ToolRecordingLog) Flush(ctx context.Context) error {
 235         r.mutex.Lock()
 236         defer r.mutex.Unlock()
 237
 238         state := "absent"
 239         _, err := os.Stat(r.marker)
 240         if err == nil {
 241                 state = "present"
 242         }
 243
 244         r.events = append(r.events, "flush, marker "+state)
 245
 246         return ctx.Err()
 247 }
 248
 249 func (r *ToolRecordingLog) Events() []string {
 250         r.mutex.Lock()
 251         defer r.mutex.Unlock()
 252
 253         return append([]string(nil), r.events...)
 254 }
 255
 256 // newTestToolbox returns a toolbox with log, an empty chat, and the working
 257 // directory dir.
 258 func newTestToolbox(t *testing.T, log ToolLog, dir string) *Toolbox {
 259         t.Helper()
 260
 261         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
 262         require.NoError(t, err)
 263
 264         t.Cleanup(func() {
 265                 require.NoError(t, memory.Close())
 266         })
 267
 268         return &Toolbox{
 269                 log:       log,
 270                 chat:      memory,
 271                 subagents: newTestSubagents(t),
 272                 workdir:   dir,
 273         }
 274 }