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

   1 package main
   2
   3 import (
   4         "bytes"
   5         "context"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "math/bits"
  11         "os"
  12         "os/exec"
  13         "path/filepath"
  14         "strconv"
  15         "strings"
  16         "syscall"
  17         "time"
  18         "unicode/utf8"
  19 )
  20
  21 const (
  22         // toolClipCharacters is the most characters of a tool output: its head
  23         // and its tail, half each.
  24         toolClipCharacters = 30000
  25
  26         // toolClipBytes is the most bytes that a tool keeps of each end of an
  27         // output, enough for toolClipCharacters/2 characters of 4 bytes.
  28         toolClipBytes = 4 * toolClipCharacters / 2
  29
  30         // toolShellWaitDelay is how long a command can keep its output open
  31         // after it ends or after its kill.
  32         toolShellWaitDelay = 2 * time.Second
  33 )
  34
  35 // ToolLog is the log that the tools write their calls and results to.
  36 type ToolLog interface {
  37         Append(kind, text string) (Message, error)
  38         Flush(ctx context.Context) error
  39 }
  40
  41 // ToolChat gives the messages and the nodes that zoom and date read.
  42 type ToolChat interface {
  43         Len() int
  44         Message(id int) (Message, error)
  45         Node(level, index int) (Node, bool, error)
  46 }
  47
  48 // Toolbox runs the tool calls of one agent. Each call and each result goes to
  49 // the log as it happens: the call as a tool message, the result as an echo
  50 // message. The call is durable before its tool runs, so the log never misses an
  51 // action outside the agent. Commands and files run as the user in the working
  52 // directory, with no sandbox. Zoom with a name reads the log of a subagent.
  53 // The toolbox of the agent has the runs of its subagents for spawn and tell.
  54 // The toolbox of a subagent has none, so spawn and tell are unknown tools
  55 // there.
  56 type Toolbox struct {
  57         log       ToolLog
  58         chat      ToolChat
  59         subagents *Subagents
  60         runs      *SubagentRuns
  61         workdir   string
  62 }
  63
  64 // Call logs the call of the tool name with its JSON arguments, makes the log
  65 // durable, runs the tool, and logs its clipped output, which it returns. A
  66 // failed tool gives its error as the output. An error of Call itself is an
  67 // error of the log. When ctx ends before the flush, the tool does not run.
  68 func (b *Toolbox) Call(ctx context.Context, name, arguments string) (string, error) {
  69         _, err := b.log.Append("tool", toolCallText(name, arguments))
  70         if err != nil {
  71                 return "", err
  72         }
  73
  74         err = b.log.Flush(ctx)
  75         if err != nil && ctx.Err() == nil {
  76                 return "", err
  77         }
  78
  79         output := "not run: the turn was interrupted"
  80         if err == nil {
  81                 output = toolClip(b.run(ctx, name, arguments))
  82         }
  83
  84         _, err = b.log.Append("echo", output)
  85         if err != nil {
  86                 return "", err
  87         }
  88
  89         return output, nil
  90 }
  91
  92 // run runs one tool and returns its output, or its error as text.
  93 func (b *Toolbox) run(ctx context.Context, name, arguments string) string {
  94         var output string
  95         var err error
  96         switch name {
  97         case "zoom":
  98                 var args ToolZoomArguments
  99                 err = toolDecodeArguments(arguments, &args)
 100                 switch {
 101                 case err != nil:
 102                 case args.Name != nil:
 103                         output, err = b.zoomAgent(ctx, *args.Name)
 104                 default:
 105                         output, err = b.zoom(args.ID, args.N)
 106                 }
 107         case "date":
 108                 var args ToolDateArguments
 109                 err = toolDecodeArguments(arguments, &args)
 110                 if err == nil {
 111                         output, err = b.date(args.ID)
 112                 }
 113         case "shell":
 114                 var args ToolShellArguments
 115                 err = toolDecodeArguments(arguments, &args)
 116                 if err == nil {
 117                         output = b.shell(ctx, args.Command)
 118                 }
 119         case "read":
 120                 var args ToolPathArguments
 121                 err = toolDecodeArguments(arguments, &args)
 122                 if err == nil {
 123                         output, err = b.read(args.Path)
 124                 }
 125         case "write":
 126                 var args ToolWriteArguments
 127                 err = toolDecodeArguments(arguments, &args)
 128                 if err == nil {
 129                         output, err = b.write(args.Path, args.Text)
 130                 }
 131         case "edit":
 132                 var args ToolEditArguments
 133                 err = toolDecodeArguments(arguments, &args)
 134                 if err == nil {
 135                         output, err = b.edit(args.Path, args.Old, args.New)
 136                 }
 137         case "spawn":
 138                 var args ToolSpawnArguments
 139                 err = toolDecodeArguments(arguments, &args)
 140                 switch {
 141                 case b.runs == nil:
 142                         err = fmt.Errorf("unknown tool %q", name)
 143                 case err == nil:
 144                         output, err = b.runs.Spawn(ctx, args.Tasks, args.Effort)
 145                 }
 146         case "tell":
 147                 var args ToolTellArguments
 148                 err = toolDecodeArguments(arguments, &args)
 149                 switch {
 150                 case b.runs == nil:
 151                         err = fmt.Errorf("unknown tool %q", name)
 152                 case err == nil:
 153                         output, err = b.runs.Tell(ctx, args.Name, args.Message)
 154                 }
 155         default:
 156                 err = fmt.Errorf("unknown tool %q", name)
 157         }
 158
 159         if err != nil {
 160                 return "error: " + err.Error()
 161         }
 162
 163         return output
 164 }
 165
 166 // zoomAgent returns the whole log of the subagent name.
 167 func (b *Toolbox) zoomAgent(ctx context.Context, name string) (string, error) {
 168         subagent, err := b.subagents.Get(ctx, name)
 169         if err != nil {
 170                 return "", err
 171         }
 172
 173         return subagent.Show()
 174 }
 175
 176 // zoom returns message id whole for n = 1, or else the two lines under the line
 177 // id+n.
 178 func (b *Toolbox) zoom(id, n int) (string, error) {
 179         if n < 1 || bits.OnesCount(uint(n)) != 1 || id < 0 || id%n != 0 {
 180                 return "", fmt.Errorf("line %d+%d does not exist: n must be a power of 2 and id a multiple of n", id, n)
 181         }
 182
 183         if id+n > b.chat.Len() {
 184                 return "", fmt.Errorf("line %d+%d goes past the last message, %d", id, n, b.chat.Len()-1)
 185         }
 186
 187         if n == 1 {
 188                 message, err := b.chat.Message(id)
 189                 if err != nil {
 190                         return "", err
 191                 }
 192
 193                 return leafSource(message), nil
 194         }
 195
 196         level := bits.TrailingZeros(uint(n)) - 1
 197         var lines []string
 198         for _, index := range []int{2 * id / n, 2*id/n + 1} {
 199                 node, built, err := b.chat.Node(level, index)
 200                 if err != nil {
 201                         return "", err
 202                 }
 203
 204                 text := viewPending
 205                 if built {
 206                         text = strings.ReplaceAll(node.Text, "\n", " ")
 207                 }
 208
 209                 lines = append(lines, fmt.Sprintf("%d+%d|%s", index<<level, 1<<level, text))
 210         }
 211
 212         return strings.Join(lines, "\n"), nil
 213 }
 214
 215 // date returns the date and time of message id, with its weekday.
 216 func (b *Toolbox) date(id int) (string, error) {
 217         if id < 0 || id >= b.chat.Len() {
 218                 return "", fmt.Errorf("message %d does not exist", id)
 219         }
 220
 221         message, err := b.chat.Message(id)
 222         if err != nil {
 223                 return "", err
 224         }
 225
 226         date, err := time.Parse(time.RFC3339Nano, message.Date)
 227         if err != nil {
 228                 return "", err
 229         }
 230
 231         return date.Local().Format("Monday 2006-01-02 15:04:05 -07:00"), nil
 232 }
 233
 234 // shell runs command with bash in the working directory and returns its output
 235 // and error output. A failed command adds its exit status. When ctx ends, the
 236 // whole process group of the command is killed.
 237 func (b *Toolbox) shell(ctx context.Context, command string) string {
 238         cmd := exec.CommandContext(ctx, "bash", "-c", command)
 239         cmd.Dir = b.workdir
 240         cmd.SysProcAttr = &syscall.SysProcAttr{
 241                 Setpgid: true,
 242         }
 243         cmd.Cancel = func() error {
 244                 return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
 245         }
 246         cmd.WaitDelay = toolShellWaitDelay
 247
 248         output := &ToolHeadTail{}
 249         cmd.Stdout = output
 250         cmd.Stderr = output
 251         err := cmd.Run()
 252
 253         text := output.String()
 254         var exit *exec.ExitError
 255         switch {
 256         case ctx.Err() != nil:
 257                 text += "\n[interrupted]"
 258         case errors.As(err, &exit):
 259                 text += "\n[" + exit.String() + "]"
 260         case err != nil:
 261                 text += "\n[error: " + err.Error() + "]"
 262         case text == "":
 263                 text = "(no output)"
 264         }
 265
 266         return text
 267 }
 268
 269 // read returns the text of the file at path.
 270 func (b *Toolbox) read(path string) (text string, err error) {
 271         file, err := os.Open(b.path(path))
 272         if err != nil {
 273                 return "", err
 274         }
 275
 276         defer func() {
 277                 err = errors.Join(err, file.Close())
 278         }()
 279
 280         output := &ToolHeadTail{}
 281         _, err = io.Copy(output, file)
 282         if err != nil {
 283                 return "", err
 284         }
 285
 286         return output.String(), nil
 287 }
 288
 289 // write replaces the file at path with text, and makes its missing
 290 // directories.
 291 func (b *Toolbox) write(path, text string) (string, error) {
 292         path = b.path(path)
 293         err := os.MkdirAll(filepath.Dir(path), 0o755)
 294         if err != nil {
 295                 return "", err
 296         }
 297
 298         err = os.WriteFile(path, []byte(text), 0o644)
 299         if err != nil {
 300                 return "", err
 301         }
 302
 303         return fmt.Sprintf("wrote %d bytes to %s", len(text), path), nil
 304 }
 305
 306 // edit replaces the one occurrence of old in the file at path with
 307 // replacement. It keeps the mode of the file.
 308 func (b *Toolbox) edit(path, old, replacement string) (string, error) {
 309         path = b.path(path)
 310         data, err := os.ReadFile(path)
 311         if err != nil {
 312                 return "", err
 313         }
 314
 315         if old == "" {
 316                 return "", errors.New("old is empty")
 317         }
 318
 319         count := strings.Count(string(data), old)
 320         if count != 1 {
 321                 return "", fmt.Errorf("old occurs %d times in %s, not once", count, path)
 322         }
 323
 324         info, err := os.Stat(path)
 325         if err != nil {
 326                 return "", err
 327         }
 328
 329         err = os.WriteFile(path, []byte(strings.Replace(string(data), old, replacement, 1)), info.Mode().Perm())
 330         if err != nil {
 331                 return "", err
 332         }
 333
 334         return "edited " + path, nil
 335 }
 336
 337 // path returns path, relative to the working directory when it is not
 338 // absolute.
 339 func (b *Toolbox) path(path string) string {
 340         if filepath.IsAbs(path) {
 341                 return path
 342         }
 343
 344         return filepath.Join(b.workdir, path)
 345 }
 346
 347 // ToolZoomArguments are the arguments of zoom. A Name that is not nil names
 348 // the subagent whose whole log zoom gives, and then ID and N are not used.
 349 type ToolZoomArguments struct {
 350         ID   int     `json:"id"`
 351         N    int     `json:"n"`
 352         Name *string `json:"name"`
 353 }
 354
 355 // ToolDateArguments are the arguments of date.
 356 type ToolDateArguments struct {
 357         ID int `json:"id"`
 358 }
 359
 360 // ToolShellArguments are the arguments of shell.
 361 type ToolShellArguments struct {
 362         Command string `json:"command"`
 363 }
 364
 365 // ToolPathArguments are the arguments of read.
 366 type ToolPathArguments struct {
 367         Path string `json:"path"`
 368 }
 369
 370 // ToolWriteArguments are the arguments of write.
 371 type ToolWriteArguments struct {
 372         Path string `json:"path"`
 373         Text string `json:"text"`
 374 }
 375
 376 // ToolEditArguments are the arguments of edit.
 377 type ToolEditArguments struct {
 378         Path string `json:"path"`
 379         Old  string `json:"old"`
 380         New  string `json:"new"`
 381 }
 382
 383 // ToolSpawnArguments are the arguments of spawn. A nil Effort is the effort of
 384 // the agent.
 385 type ToolSpawnArguments struct {
 386         Tasks  []string `json:"tasks"`
 387         Effort *string  `json:"effort"`
 388 }
 389
 390 // ToolTellArguments are the arguments of tell.
 391 type ToolTellArguments struct {
 392         Name    string `json:"name"`
 393         Message string `json:"message"`
 394 }
 395
 396 // ToolHeadTail keeps the first and the last toolClipBytes bytes written to it,
 397 // and counts the bytes between them that it drops.
 398 type ToolHeadTail struct {
 399         head    []byte
 400         tail    []byte
 401         dropped int
 402 }
 403
 404 // String returns the output, clipped when it is too long. Invalid UTF-8 becomes
 405 // the replacement character.
 406 func (h *ToolHeadTail) String() string {
 407         replacement := string(utf8.RuneError)
 408         if h.dropped == 0 {
 409                 return toolClip(strings.ToValidUTF8(string(h.head)+string(h.tail), replacement))
 410         }
 411
 412         head := toolFirstCharacters(strings.ToValidUTF8(string(h.head), replacement), toolClipCharacters/2)
 413         tail := toolLastCharacters(strings.ToValidUTF8(string(h.tail), replacement), toolClipCharacters/2)
 414         cut := len(h.head) + h.dropped + len(h.tail) - len(head) - len(tail)
 415
 416         return toolJoinClipped(head, tail, cut)
 417 }
 418
 419 // Write keeps the bytes of p at the ends of the output.
 420 func (h *ToolHeadTail) Write(p []byte) (int, error) {
 421         n := len(p)
 422         room := min(toolClipBytes-len(h.head), len(p))
 423         h.head = append(h.head, p[:room]...)
 424         p = p[room:]
 425
 426         h.tail = append(h.tail, p...)
 427         if len(h.tail) > toolClipBytes {
 428                 cut := len(h.tail) - toolClipBytes
 429                 h.dropped += cut
 430                 h.tail = bytes.Clone(h.tail[cut:])
 431         }
 432
 433         return n, nil
 434 }
 435
 436 // toolCallText returns the log text of a tool call: the name and the JSON
 437 // arguments.
 438 func toolCallText(name, arguments string) string {
 439         return name + " " + arguments
 440 }
 441
 442 // toolClip returns text when it has at most toolClipCharacters characters.
 443 // Otherwise it returns the first and the last toolClipCharacters/2 characters,
 444 // with a mark of the cut between them.
 445 func toolClip(text string) string {
 446         if utf8.RuneCountInString(text) <= toolClipCharacters {
 447                 return text
 448         }
 449
 450         head := toolFirstCharacters(text, toolClipCharacters/2)
 451         tail := toolLastCharacters(text, toolClipCharacters/2)
 452
 453         return toolJoinClipped(head, tail, len(text)-len(head)-len(tail))
 454 }
 455
 456 // toolDecodeArguments decodes the JSON arguments of a call into args. An
 457 // unknown field is an error.
 458 func toolDecodeArguments(arguments string, args any) error {
 459         decoder := json.NewDecoder(strings.NewReader(arguments))
 460         decoder.DisallowUnknownFields()
 461         err := decoder.Decode(args)
 462         if err != nil {
 463                 return fmt.Errorf("arguments: %w", err)
 464         }
 465
 466         return nil
 467 }
 468
 469 // toolFirstCharacters returns the first count characters of text.
 470 func toolFirstCharacters(text string, count int) string {
 471         end := 0
 472         for range count {
 473                 if end >= len(text) {
 474                         break
 475                 }
 476
 477                 _, size := utf8.DecodeRuneInString(text[end:])
 478                 end += size
 479         }
 480
 481         return text[:end]
 482 }
 483
 484 // toolLastCharacters returns the last count characters of text.
 485 func toolLastCharacters(text string, count int) string {
 486         start := len(text)
 487         for range count {
 488                 if start <= 0 {
 489                         break
 490                 }
 491
 492                 _, size := utf8.DecodeLastRuneInString(text[:start])
 493                 start -= size
 494         }
 495
 496         return text[start:]
 497 }
 498
 499 // toolJoinClipped joins the head and the tail of a clipped output with a mark
 500 // of the cut bytes.
 501 func toolJoinClipped(head, tail string, cut int) string {
 502         return head + "\n[... " + strconv.Itoa(cut) + " bytes cut ...]\n" + tail
 503 }