]> git.giacomolli.uk - onechat/blob - tools_test.go
Talk to the agent through a Telegram bot
[onechat] / 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 := &recordingLog{
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 := &recordingLog{
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, &recordingLog{}, 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, &recordingLog{}, 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, &recordingLog{}, 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", clipCharacters)
174 require.Equal(t, short, clip(short))
175
176 text := strings.Repeat("h", clipCharacters/2) + strings.Repeat("é", 5000) + strings.Repeat("t", clipCharacters/2)
177 require.Equal(t, strings.Repeat("h", clipCharacters/2)+"\n[... 10000 bytes cut ...]\n"+strings.Repeat("t", clipCharacters/2), clip(text))
178
179 half := strings.Repeat("é", clipCharacters/2)
180 clipped := clip(strings.Repeat("é", clipCharacters+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 := &headTail{}
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-clipCharacters))
196 require.Equal(t, clipCharacters+len(fmt.Sprintf("\n[... %d bytes cut ...]\n", 1000*1000-clipCharacters)), len(text))
197 }
198
199 // recordingLog is a chat log that records each append and flush. The flush
200 // records whether the marker file exists.
201 type recordingLog struct {
202 mutex sync.Mutex
203 marker string
204 events []string
205 }
206
207 func (r *recordingLog) 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 *recordingLog) 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 *recordingLog) 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 chatLog, 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 }