byre/internal/agents/files.go

package agents

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
	"unicode/utf8"

	"g.fc.run/ai/byre/internal/grep"
	"g.fc.run/ai/byre/internal/mcp"
)

const grepDescription = "Search files for lines matching a regular expression (Go RE2 syntax), as ripgrep does: a directory is searched recursively, skipping .git, what .gitignore files exclude, and binary files. " +
	"Each match comes back as path:line:text, at most 64 KiB in all."

func (m *Manager) fileTools(l *live) []mcp.Tool {
	return []mcp.Tool{
		{
			Name: "read",
			Description: "Read a file on the Mac byre runs on; byre reads it itself, outside any sandbox your own tools have. Text comes back with line numbers, starting at line 1 (or offset), at most limit lines. " +
				"A PNG, JPEG, GIF or WebP image comes back as the image. Other binary files fail.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"path": {"type": "string", "description": "The file, relative to your worktree, absolute, or starting with ~/."},
					"offset": {"type": "integer", "description": "The line number to start from (default 1)."},
					"limit": {"type": "integer", "description": "The most lines to return (default 2000)."}
				},
				"required": ["path"]
			}`),
			Rich: func(ctx context.Context, raw json.RawMessage) mcp.Result {
				return m.read(l, raw)
			},
		},
		{
			Name:        "write",
			Description: "Write a file on the Mac byre runs on, replacing its contents, creating it and its parent directories when missing; byre writes it itself, outside any sandbox your own tools have.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"path": {"type": "string", "description": "The file, relative to your worktree, absolute, or starting with ~/."},
					"content": {"type": "string"}
				},
				"required": ["path", "content"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.write(l, raw)
			},
		},
		{
			Name: "edit",
			Description: "Replace old_string with new_string in a text file on the Mac byre runs on; byre edits it itself, outside any sandbox your own tools have. " +
				"old_string must match the file exactly, whitespace included, and appear exactly once; otherwise the file is left as it was.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"path": {"type": "string", "description": "The file, relative to your worktree, absolute, or starting with ~/."},
					"old_string": {"type": "string"},
					"new_string": {"type": "string"}
				},
				"required": ["path", "old_string", "new_string"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.edit(l, raw)
			},
		},
		{
			Name:        "grep",
			Description: grepDescription + " It searches the Mac byre runs on, outside any sandbox your own tools have.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"pattern": {"type": "string"},
					"path": {"type": "string", "description": "The file or directory to search, relative to your worktree, absolute, or starting with ~/ (default your worktree)."},
					"glob": {"type": "string", "description": "Search only files whose names match this glob, e.g. *.go."},
					"ignore_case": {"type": "boolean"},
					"files_only": {"type": "boolean", "description": "List the matching files instead of their lines."}
				},
				"required": ["pattern"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.grep(ctx, l, raw)
			},
		},
	}
}

func (l *live) resolve(path string) (string, error) {
	if path == "" {
		return "", errors.New("path is empty")
	}
	rest, ok := strings.CutPrefix(path, "~/")
	if ok {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", err
		}
		return filepath.Join(home, rest), nil
	}
	if filepath.IsAbs(path) {
		return path, nil
	}
	return filepath.Join(l.dir, path), nil
}

const readMax = 64 << 20

func regular(full string) (os.FileInfo, error) {
	info, err := os.Stat(full)
	if err != nil {
		return nil, err
	}
	if !info.Mode().IsRegular() {
		return nil, fmt.Errorf("%s isn't a regular file", full)
	}
	return info, nil
}

func readLocal(full string) ([]byte, error) {
	info, err := regular(full)
	if err != nil {
		return nil, err
	}
	if info.Size() > readMax {
		return nil, fmt.Errorf("%s is %d bytes, over the 64 MiB limit", full, info.Size())
	}
	return os.ReadFile(full)
}

type readArgs struct {
	Path   string `json:"path"`
	Offset *int   `json:"offset"`
	Limit  *int   `json:"limit"`
}

func (m *Manager) read(l *live, raw json.RawMessage) mcp.Result {
	args := readArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return mcp.TextResult(fmt.Sprintf("malformed read call: %s", err), true)
	}
	full, err := l.resolve(args.Path)
	if err != nil {
		return mcp.TextResult(err.Error(), true)
	}
	data, err := readLocal(full)
	if err != nil {
		return mcp.TextResult(err.Error(), true)
	}
	return readResult(args.Path, data, args.Offset, args.Limit)
}

type writeArgs struct {
	Path    string `json:"path"`
	Content string `json:"content"`
}

func (m *Manager) write(l *live, raw json.RawMessage) (string, error) {
	args := writeArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed write call: %w", err)
	}
	full, err := l.resolve(args.Path)
	if err != nil {
		return "", err
	}
	_, err = regular(full)
	if err != nil && !errors.Is(err, fs.ErrNotExist) {
		return "", err
	}
	err = os.MkdirAll(filepath.Dir(full), 0o755)
	if err != nil {
		return "", err
	}
	err = os.WriteFile(full, []byte(args.Content), 0o644)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Wrote %d bytes to %s.", len(args.Content), args.Path), nil
}

type editArgs struct {
	Path      string `json:"path"`
	OldString string `json:"old_string"`
	NewString string `json:"new_string"`
}

func (m *Manager) edit(l *live, raw json.RawMessage) (string, error) {
	args := editArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed edit call: %w", err)
	}
	full, err := l.resolve(args.Path)
	if err != nil {
		return "", err
	}
	data, err := readLocal(full)
	if err != nil {
		return "", err
	}
	edited, err := replaceOnce(args.Path, data, args.OldString, args.NewString)
	if err != nil {
		return "", err
	}
	err = os.WriteFile(full, edited, 0o644)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Edited %s.", args.Path), nil
}

type grepArgs struct {
	Pattern    string `json:"pattern"`
	Path       string `json:"path"`
	Glob       string `json:"glob"`
	IgnoreCase bool   `json:"ignore_case"`
	FilesOnly  bool   `json:"files_only"`
}

func (a grepArgs) options() grep.Options {
	return grep.Options{Glob: a.Glob, IgnoreCase: a.IgnoreCase, FilesOnly: a.FilesOnly}
}

func (m *Manager) grep(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := grepArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed grep call: %w", err)
	}
	path := args.Path
	if path == "" {
		path = "."
	}
	if strings.HasPrefix(path, "~/") {
		path, err = l.resolve(path)
		if err != nil {
			return "", err
		}
	}
	cctx, cancel := context.WithTimeout(ctx, fileTimeout)
	defer cancel()
	return grep.Search(cctx, args.Pattern, l.dir, path, args.options())
}

var imageTypes = map[string]string{".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif", ".webp": "image/webp"}

func readResult(path string, data []byte, offset, limit *int) mcp.Result {
	mime, ok := imageTypes[strings.ToLower(filepath.Ext(path))]
	if ok {
		return mcp.Result{Content: []mcp.Content{{Type: "image", Data: base64.StdEncoding.EncodeToString(data), MimeType: mime}}}
	}
	text, err := numberLines(path, data, offset, limit)
	if err != nil {
		return mcp.TextResult(err.Error(), true)
	}
	return mcp.TextResult(text, false)
}

func numberLines(path string, data []byte, offset, limit *int) (string, error) {
	if !utf8.Valid(data) {
		return "", fmt.Errorf("%s isn't text", path)
	}
	if len(data) == 0 {
		return fmt.Sprintf("%s is empty.", path), nil
	}
	start := 1
	if offset != nil {
		start = *offset
	}
	n := defaultReadLimit
	if limit != nil {
		n = *limit
	}
	if start < 1 || n < 1 {
		return "", errors.New("offset and limit must be at least 1")
	}
	lines := strings.Split(strings.TrimSuffix(string(data), "\n"), "\n")
	if start > len(lines) {
		return "", fmt.Errorf("%s has %d lines; offset %d is past the end", path, len(lines), start)
	}
	end := min(start-1+n, len(lines))
	b := &strings.Builder{}
	for i := start - 1; i < end; i++ {
		fmt.Fprintf(b, "%6d\t%s\n", i+1, lines[i])
	}
	if end < len(lines) {
		fmt.Fprintf(b, "(%d more lines; read on with offset %d)\n", len(lines)-end, end+1)
	}
	return b.String(), nil
}

func replaceOnce(path string, data []byte, old, new string) ([]byte, error) {
	if old == "" {
		return nil, errors.New("old_string is empty")
	}
	if old == new {
		return nil, errors.New("old_string and new_string are the same")
	}
	if !utf8.Valid(data) {
		return nil, fmt.Errorf("%s isn't text", path)
	}
	text := string(data)
	n := strings.Count(text, old)
	if n == 0 {
		return nil, fmt.Errorf("old_string isn't in %s", path)
	}
	if n > 1 {
		return nil, fmt.Errorf("old_string appears %d times in %s; include more of the text around it so it appears once", n, path)
	}
	return []byte(strings.Replace(text, old, new, 1)), nil
}