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