Separate multiline Go statements

[onechat] / tools.go

package main

import (
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "math/bits"
        "os"
        "os/exec"
        "path/filepath"
        "strconv"
        "strings"
        "syscall"
        "time"
        "unicode/utf8"
)

const (
        // toolClipCharacters is the most characters of a tool output: its head
        // and its tail, half each.
        toolClipCharacters = 30000

        // toolClipBytes is the most bytes that a tool keeps of each end of an
        // output, enough for toolClipCharacters/2 characters of 4 bytes.
        toolClipBytes = 4 * toolClipCharacters / 2

        // toolShellOutputTimeout is how long a command can keep its output open
        // after it ends or after its kill.
        toolShellOutputTimeout = 2 * time.Second
)

// ToolLog is the log that the tools write their calls and results to.
type ToolLog interface {
        Append(kind, text string) (Message, error)
        Flush(ctx context.Context) error
}

// ToolChat gives the messages and the nodes that zoom and date read.
type ToolChat interface {
        Len() int
        Message(id int) (Message, error)
        Node(level, index int) (Node, bool, error)
}

// Toolbox runs the tool calls of one agent. Each call and each result goes to
// the log as it happens: the call as a tool message, the result as an echo
// message. The call is durable before its tool runs, so the log never misses an
// action outside the agent. Commands and files run as the user in the working
// directory, with no sandbox. Zoom with a name reads a child's view or a
// subagent's log. The toolbox of the agent has the runs of its subagents for
// spawn and tell, the WhatsApp link for the WhatsApp tools, and the schedules
// for the schedule tools. Each permanent agent has the parent-child queue. Only
// the root has child creation and source routes. A child has neither of them
// nor a WhatsApp link. Only the root can show images to the user. The toolbox
// of a subagent has none of them, so these are unknown tools there.
type Toolbox struct {
        log       ToolLog
        chat      ToolChat
        subagents *Subagents
        runs      *SubagentRuns
        whatsapp  *WhatsApp
        schedules *Schedules
        workdir   string
        agentName string
        messages  *DaemonMessages
        children  *DaemonChildren
        routes    *Routes
        images    *ImageStore
        imageShow *ImageShow
}

// ToolZoomArguments are the arguments of zoom. A Name that is not nil names the
// child whose view or the subagent whose whole log zoom gives. ID and N are
// then not used.
type ToolZoomArguments struct {
        ID   int     `json:"id"`
        N    int     `json:"n"`
        Name *string `json:"name"`
}

// ToolDateArguments are the arguments of date.
type ToolDateArguments struct {
        ID int `json:"id"`
}

// ToolShellArguments are the arguments of shell.
type ToolShellArguments struct {
        Command string `json:"command"`
}

// ToolPathArguments are the arguments of read.
type ToolPathArguments struct {
        Path string `json:"path"`
}

// ToolWriteArguments are the arguments of write.
type ToolWriteArguments struct {
        Path string `json:"path"`
        Text string `json:"text"`
}

// ToolEditArguments are the arguments of edit.
type ToolEditArguments struct {
        Path string `json:"path"`
        Old  string `json:"old"`
        New  string `json:"new"`
}

// ToolSpawnArguments are the arguments of spawn. A nil Effort is the effort of
// the agent.
type ToolSpawnArguments struct {
        Tasks  []string `json:"tasks"`
        Effort *string  `json:"effort"`
}

// ToolAgentCreateArguments are the arguments of agent_create.
type ToolAgentCreateArguments struct {
        Name           string `json:"name"`
        Responsibility string `json:"responsibility"`
}

// ToolRouteArguments are the arguments of route.
type ToolRouteArguments struct {
        Source string `json:"source"`
        Agent  string `json:"agent"`
}

// ToolRouteSourceArguments are the arguments of unroute.
type ToolRouteSourceArguments struct {
        Source string `json:"source"`
}

// ToolTellArguments are the arguments of tell.
type ToolTellArguments struct {
        Name    string `json:"name"`
        Message string `json:"message"`
}

// ToolImageShowArguments are the arguments of image_show.
type ToolImageShowArguments struct {
        Image   string `json:"image"`
        Caption string `json:"caption"`
}

// ToolWhatsAppFindArguments are the arguments of whatsapp_find.
type ToolWhatsAppFindArguments struct {
        Name string `json:"name"`
}

// ToolWhatsAppSendArguments are the arguments of whatsapp_send.
type ToolWhatsAppSendArguments struct {
        ID    string `json:"id"`
        Chat  string `json:"chat"`
        Text  string `json:"text"`
        Image string `json:"image"`
}

// ToolScheduleSetArguments are the arguments of schedule_set. Exactly one of At
// and Cron is not nil.
type ToolScheduleSetArguments struct {
        Text string  `json:"text"`
        At   *string `json:"at"`
        Cron *string `json:"cron"`
}

// ToolScheduleRemoveArguments are the arguments of schedule_remove.
type ToolScheduleRemoveArguments struct {
        ID int `json:"id"`
}

// ToolHeadTail keeps the first and the last toolClipBytes bytes written to it,
// and counts the bytes between them that it drops.
type ToolHeadTail struct {
        head    []byte
        tail    []byte
        dropped int
}

// Call logs the call of the tool name with its JSON arguments, makes the log
// durable, runs the tool, and logs its clipped text. It returns that text or
// the text and image parts of a single-message zoom. Pixels never enter the
// log. A failed tool gives its error as the output. An error of Call itself is
// an error of the log. When ctx ends before the flush, the tool does not run.
func (b *Toolbox) Call(ctx context.Context, name, arguments string) (any, error) {
        _, err := b.log.Append("tool", toolCallText(name, arguments))
        if err != nil {
                return "", fmt.Errorf("append tool call: %w", err)
        }

        err = b.log.Flush(ctx)
        if err != nil && ctx.Err() == nil {
                return "", fmt.Errorf("flush tool call: %w", err)
        }

        output := "not run: the turn was interrupted"
        if err == nil {
                output = toolClip(b.run(ctx, name, arguments))
        }

        var (
                result    any = output
                resultErr error
        )

        if err == nil {
                result, resultErr = imageToolOutput(b, name, arguments, output)
        }

        if resultErr != nil {
                output = "error: " + resultErr.Error()
                result = output
        }

        _, err = b.log.Append("echo", output)
        if err != nil {
                return "", fmt.Errorf("append tool result: %w", err)
        }

        return result, nil
}

// run runs one tool and returns its output, or its error as text.
func (b *Toolbox) run(ctx context.Context, name, arguments string) string {
        var (
                output string
                err    error
        )

        switch name {
        case "zoom":
                var args ToolZoomArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case err != nil:
                case args.Name != nil:
                        output, err = b.zoomAgent(ctx, *args.Name)
                default:
                        output, err = b.zoom(args.ID, args.N)
                }
        case "date":
                var args ToolDateArguments

                err = toolDecodeArguments(arguments, &args)
                if err == nil {
                        output, err = b.date(args.ID)
                }
        case "shell":
                var args ToolShellArguments

                err = toolDecodeArguments(arguments, &args)
                if err == nil {
                        output = b.shell(ctx, args.Command)
                }
        case "read":
                var args ToolPathArguments

                err = toolDecodeArguments(arguments, &args)
                if err == nil {
                        output, err = b.read(args.Path)
                }
        case "write":
                var args ToolWriteArguments

                err = toolDecodeArguments(arguments, &args)
                if err == nil {
                        output, err = b.write(args.Path, args.Text)
                }
        case "edit":
                var args ToolEditArguments

                err = toolDecodeArguments(arguments, &args)
                if err == nil {
                        output, err = b.edit(args.Path, args.Old, args.New)
                }
        case "spawn":
                var args ToolSpawnArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.runs == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.runs.Spawn(ctx, args.Tasks, args.Effort)
                }
        case "agent_create":
                var args ToolAgentCreateArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.children == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.children.Create(args.Name, args.Responsibility)
                }
        case "route", "unroute", "routes":
                var args ToolRouteArguments

                switch name {
                case "routes":
                        err = toolDecodeArguments(arguments, &struct{}{})
                case "unroute":
                        var source ToolRouteSourceArguments

                        err = toolDecodeArguments(arguments, &source)
                        args.Source = source.Source
                default:
                        err = toolDecodeArguments(arguments, &args)
                }

                switch {
                case b.routes == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err != nil:
                case name == "routes":
                        output = b.routes.List()
                case name == "route" && (b.messages == nil || b.messages.Child(b.agentName, args.Agent) == nil):
                        err = errors.New("the agent must be a child of the root")
                default:
                        err = b.routes.Set(args.Source, args.Agent)
                        output = "Route removed."
                        if name == "route" {
                                output = "Route set."
                        }
                }
        case "tell":
                var args ToolTellArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.runs == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil && b.messages != nil && b.messages.Child(b.agentName, args.Name) != nil:
                        err = b.messages.Tell(b.agentName, args.Name, args.Message)
                        output = "Sent. " + args.Name + "'s reply will reach you as a message."
                case err == nil:
                        output, err = b.runs.Tell(ctx, args.Name, args.Message)
                }
        case "image_show":
                var args ToolImageShowArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.imageShow == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.imageShow.Show(ctx, args.Image, args.Caption)
                }
        case "whatsapp_find":
                var args ToolWhatsAppFindArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.whatsapp == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.whatsapp.Find(ctx, args.Name)
                }
        case "whatsapp_send":
                var args ToolWhatsAppSendArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.whatsapp == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.whatsapp.Send(ctx, args.ID, args.Chat, args.Text, args.Image)
                }
        case "schedule_set":
                var args ToolScheduleSetArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.schedules == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.schedules.Set(args.Text, args.At, args.Cron)
                }
        case "schedule_list":
                err = toolDecodeArguments(arguments, &struct{}{})
                switch {
                case b.schedules == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output = b.schedules.List()
                }
        case "schedule_remove":
                var args ToolScheduleRemoveArguments

                err = toolDecodeArguments(arguments, &args)
                switch {
                case b.schedules == nil:
                        err = fmt.Errorf("unknown tool %q", name)
                case err == nil:
                        output, err = b.schedules.Remove(args.ID)
                }
        default:
                err = fmt.Errorf("unknown tool %q", name)
        }

        if err != nil {
                return "error: " + err.Error()
        }

        return output
}

// zoomAgent returns a child's view or the whole log of a subagent.
func (b *Toolbox) zoomAgent(ctx context.Context, name string) (string, error) {
        if b.messages != nil {
                child := b.messages.Child(b.agentName, name)
                if child != nil {
                        return strings.Join(child.memory.ViewLines(nil), ""), nil
                }
        }

        subagent, err := b.subagents.Get(ctx, name)
        if err != nil {
                return "", err
        }

        return subagent.Show()
}

// zoom returns message id whole for n = 1, or else the two lines under the line
// id+n.
func (b *Toolbox) zoom(id, n int) (string, error) {
        if n < 1 || bits.OnesCount(uint(n)) != 1 || id < 0 || id%n != 0 {
                return "", fmt.Errorf("line %d+%d does not exist: n must be a power of 2 and id a multiple of n", id, n)
        }

        if id+n > b.chat.Len() {
                return "", fmt.Errorf("line %d+%d goes past the last message, %d", id, n, b.chat.Len()-1)
        }

        if n == 1 {
                message, err := b.chat.Message(id)
                if err != nil {
                        return "", fmt.Errorf("read zoom message: %w", err)
                }

                return leafSource(message), nil
        }

        level := bits.TrailingZeros(uint(n)) - 1

        var lines []string

        for _, index := range []int{2 * id / n, 2*id/n + 1} {
                node, built, err := b.chat.Node(level, index)
                if err != nil {
                        return "", fmt.Errorf("read zoom node: %w", err)
                }

                text := viewPending
                if built {
                        text = strings.ReplaceAll(node.Text, "\n", " ")
                }

                lines = append(lines, fmt.Sprintf("%d+%d|%s", index<<level, 1<<level, text))
        }

        return strings.Join(lines, "\n"), nil
}

// date returns the date and time of message id, with its weekday.
func (b *Toolbox) date(id int) (string, error) {
        if id < 0 || id >= b.chat.Len() {
                return "", fmt.Errorf("message %d does not exist", id)
        }

        message, err := b.chat.Message(id)
        if err != nil {
                return "", fmt.Errorf("read dated message: %w", err)
        }

        date, err := time.Parse(time.RFC3339Nano, message.Date)
        if err != nil {
                return "", fmt.Errorf("parse message date: %w", err)
        }

        return date.Local().Format("Monday 2006-01-02 15:04:05 -07:00"), nil
}

// shell runs command with bash in the working directory and returns its output
// and error output. A failed command adds its exit status. When ctx ends, the
// whole process group of the command is killed.
func (b *Toolbox) shell(ctx context.Context, command string) string {
        cmd := exec.CommandContext(ctx, "bash", "-c", command)
        cmd.Dir = b.workdir
        cmd.SysProcAttr = &syscall.SysProcAttr{
                Setpgid: true,
        }

        cmd.Cancel = func() error {
                return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
        }

        cmd.WaitDelay = toolShellOutputTimeout

        output := &ToolHeadTail{}
        cmd.Stdout = output
        cmd.Stderr = output
        err := cmd.Run()

        text := output.String()

        var exit *exec.ExitError

        switch {
        case ctx.Err() != nil:
                text += "\n[interrupted]"
        case errors.As(err, &exit):
                text += "\n[" + exit.String() + "]"
        case err != nil:
                text += "\n[error: " + err.Error() + "]"
        case text == "":
                text = "(no output)"
        }

        return text
}

// read returns the text of the file at path.
func (b *Toolbox) read(path string) (text string, err error) {
        file, err := os.Open(b.path(path))
        if err != nil {
                return "", fmt.Errorf("open tool file: %w", err)
        }

        defer func() {
                err = errors.Join(err, file.Close())
        }()

        output := &ToolHeadTail{}
        _, err = io.Copy(output, file)
        if err != nil {
                return "", fmt.Errorf("read tool file: %w", err)
        }

        return output.String(), nil
}

// write replaces the file at path with text, and makes its missing
// directories.
func (b *Toolbox) write(path, text string) (string, error) {
        path = b.path(path)
        err := os.MkdirAll(filepath.Dir(path), 0o755)
        if err != nil {
                return "", fmt.Errorf("create tool file directory: %w", err)
        }

        err = os.WriteFile(path, []byte(text), 0o644)
        if err != nil {
                return "", fmt.Errorf("write tool file: %w", err)
        }

        return fmt.Sprintf("wrote %d bytes to %s", len(text), path), nil
}

// edit replaces the one occurrence of old in the file at path with
// replacement. It keeps the mode of the file.
func (b *Toolbox) edit(path, old, replacement string) (string, error) {
        path = b.path(path)
        data, err := os.ReadFile(path)
        if err != nil {
                return "", fmt.Errorf("read file to edit: %w", err)
        }

        if old == "" {
                return "", errors.New("old is empty")
        }

        count := strings.Count(string(data), old)
        if count != 1 {
                return "", fmt.Errorf("old occurs %d times in %s, not once", count, path)
        }

        info, err := os.Stat(path)
        if err != nil {
                return "", fmt.Errorf("read file mode: %w", err)
        }

        err = os.WriteFile(path, []byte(strings.Replace(string(data), old, replacement, 1)), info.Mode().Perm())
        if err != nil {
                return "", fmt.Errorf("write edited file: %w", err)
        }

        return "edited " + path, nil
}

// path returns path, relative to the working directory when it is not
// absolute.
func (b *Toolbox) path(path string) string {
        if filepath.IsAbs(path) {
                return path
        }

        return filepath.Join(b.workdir, path)
}

// String returns the output, clipped when it is too long. Invalid UTF-8 becomes
// the replacement character.
func (h *ToolHeadTail) String() string {
        replacement := string(utf8.RuneError)
        if h.dropped == 0 {
                return toolClip(strings.ToValidUTF8(string(h.head)+string(h.tail), replacement))
        }

        head := toolFirstCharacters(strings.ToValidUTF8(string(h.head), replacement), toolClipCharacters/2)
        tail := toolLastCharacters(strings.ToValidUTF8(string(h.tail), replacement), toolClipCharacters/2)
        cut := len(h.head) + h.dropped + len(h.tail) - len(head) - len(tail)

        return toolJoinClipped(head, tail, cut)
}

// Write keeps the bytes of p at the ends of the output.
func (h *ToolHeadTail) Write(p []byte) (int, error) {
        n := len(p)
        room := min(toolClipBytes-len(h.head), len(p))
        h.head = append(h.head, p[:room]...)
        p = p[room:]

        h.tail = append(h.tail, p...)
        if len(h.tail) > toolClipBytes {
                cut := len(h.tail) - toolClipBytes
                h.dropped += cut
                h.tail = bytes.Clone(h.tail[cut:])
        }

        return n, nil
}

// toolCallText returns the log text of a tool call: the name and the JSON
// arguments.
func toolCallText(name, arguments string) string {
        return name + " " + arguments
}

// toolClip returns text when it has at most toolClipCharacters characters.
// Otherwise it returns the first and the last toolClipCharacters/2 characters,
// with a mark of the cut between them.
func toolClip(text string) string {
        if utf8.RuneCountInString(text) <= toolClipCharacters {
                return text
        }

        head := toolFirstCharacters(text, toolClipCharacters/2)
        tail := toolLastCharacters(text, toolClipCharacters/2)

        return toolJoinClipped(head, tail, len(text)-len(head)-len(tail))
}

// toolDecodeArguments decodes the JSON arguments of a call into args. An
// unknown field is an error.
func toolDecodeArguments(arguments string, args any) error {
        decoder := json.NewDecoder(strings.NewReader(arguments))
        decoder.DisallowUnknownFields()

        err := decoder.Decode(args)
        if err != nil {
                return fmt.Errorf("arguments: %w", err)
        }

        return nil
}

// toolFirstCharacters returns the first count characters of text.
func toolFirstCharacters(text string, count int) string {
        end := 0
        for range count {
                if end >= len(text) {
                        break
                }

                _, size := utf8.DecodeRuneInString(text[end:])
                end += size
        }

        return text[:end]
}

// toolLastCharacters returns the last count characters of text.
func toolLastCharacters(text string, count int) string {
        start := len(text)
        for range count {
                if start <= 0 {
                        break
                }

                _, size := utf8.DecodeLastRuneInString(text[:start])
                start -= size
        }

        return text[start:]
}

// toolJoinClipped joins the head and the tail of a clipped output with a mark
// of the cut bytes.
func toolJoinClipped(head, tail string, cut int) string {
        return head + "\n[... " + strconv.Itoa(cut) + " bytes cut ...]\n" + tail
}