byre/internal/agents/clients.go

package agents

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

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/clients"
	"g.fc.run/misc/byre/internal/mcp"
)

const (
	defaultBashTimeout = 2 * time.Minute
	maxBashTimeout     = 10 * time.Minute
	replyGrace         = 30 * time.Second
	copyTimeout        = 10 * time.Minute
	fileTimeout        = 2 * time.Minute
	handsTimeout       = 11 * time.Minute
	defaultReadLimit   = 2000
)

func (m *Manager) clientTools(l *live) []mcp.Tool {
	return []mcp.Tool{
		{
			Name: "list_clients",
			Description: "List the Macs with byre's app open: the human's own machines, not the one byre (and your worktree) runs on. " +
				"Each has its name, which the other client tools take as client, the human using it, and when the human last did something in byre from it. " +
				"The most recently active one is listed first; it is usually the Mac the human sits at.",
			InputSchema: json.RawMessage(`{"type": "object", "properties": {}}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.listClients(ctx)
			},
		},
		{
			Name: "client_bash",
			Description: "Run a shell command on a client Mac, as the human, with zsh as a login shell in their home directory, and get its exit status and output (stdout and stderr together, the last 64 KiB). " +
				"It runs at once, with no approval. The command runs in its own process group, which is killed at the timeout. " +
				"A command that quits byre's app on that Mac, for example to install a new build, ends this call with the client disconnecting, but the command keeps running.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"command": {"type": "string"},
					"timeout_ms": {"type": "integer", "description": "Kill the command after this long (default 120000, at most 600000)."}
				},
				"required": ["client", "command"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.clientBash(ctx, raw)
			},
		},
		{
			Name: "client_upload",
			Description: "Copy a file or directory from your worktree to a client Mac, replacing whatever is at dest there. " +
				"Directories, app bundles included, are copied whole with their permissions and symlinks.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"path": {"type": "string", "description": "The file or directory, relative to your worktree."},
					"dest": {"type": "string", "description": "Its new path on the client, absolute or starting with ~/, e.g. ~/Applications/byre.app."}
				},
				"required": ["client", "path", "dest"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.clientUpload(ctx, l, raw)
			},
		},
		{
			Name: "client_download",
			Description: "Copy a file or directory from a client Mac into your worktree, replacing whatever is at dest there. " +
				"Directories are copied whole with their permissions and symlinks.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"path": {"type": "string", "description": "The file or directory on the client, absolute or starting with ~/."},
					"dest": {"type": "string", "description": "Its new path, relative to your worktree; parent directories are created."}
				},
				"required": ["client", "path", "dest"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.clientDownload(ctx, l, raw)
			},
		},
		{
			Name: "client_read",
			Description: "Read a file on a client Mac. 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; use client_download for them.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"path": {"type": "string", "description": "The file, 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": ["client", "path"]
			}`),
			Rich: func(ctx context.Context, raw json.RawMessage) mcp.Result {
				return m.clientRead(ctx, raw)
			},
		},
		{
			Name:        "client_write",
			Description: "Write a file on a client Mac, replacing its contents, creating it and its parent directories when missing.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"path": {"type": "string", "description": "The file, absolute or starting with ~/."},
					"content": {"type": "string"}
				},
				"required": ["client", "path", "content"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.clientWrite(ctx, raw)
			},
		},
		{
			Name: "client_edit",
			Description: "Replace old_string with new_string in a text file on a client Mac. " +
				"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": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"path": {"type": "string", "description": "The file, absolute or starting with ~/."},
					"old_string": {"type": "string"},
					"new_string": {"type": "string"}
				},
				"required": ["client", "path", "old_string", "new_string"]
			}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				return m.clientEdit(ctx, raw)
			},
		},
		{
			Name: "client_hands",
			Description: "Call one of your hands tools (screen_list, capture, ax_tree, act, browser) on a client Mac instead of the Mac byre runs on: " +
				"tool is the tool's name and arguments exactly what that tool takes. Window ids, refs and displays are that Mac's, so start with screen_list there. " +
				"Every agent using that Mac shares one set of refs, so read the tree again if another agent may have. " +
				"A call fails until the human has given byre's app Screen Recording and Accessibility (and Automation of Orion, for browser) on that Mac; say so when it does.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients."},
					"tool": {"type": "string", "enum": ["screen_list", "capture", "ax_tree", "act", "browser"]},
					"arguments": {"type": "object", "description": "The tool's arguments."}
				},
				"required": ["client", "tool"]
			}`),
			Rich: func(ctx context.Context, raw json.RawMessage) mcp.Result {
				return m.clientHands(ctx, raw)
			},
		},
	}
}

func (m *Manager) listClients(ctx context.Context) (string, error) {
	status, err := m.status(ctx)
	if err != nil {
		return "", err
	}
	type row struct {
		c    clients.Client
		last *time.Time
	}
	rows := []row{}
	for _, c := range m.clients.List() {
		r := row{c: c}
		for _, s := range status.Clients {
			if s.Login == c.Login && s.Host == c.Host && s.Commit == c.Commit {
				r.last = s.LastActionAt
			}
		}
		rows = append(rows, r)
	}
	if len(rows) == 0 {
		return "No client Macs are connected.", nil
	}
	sort.SliceStable(rows, func(i, j int) bool {
		a, b := rows[i].last, rows[j].last
		return a != nil && (b == nil || a.After(*b))
	})

	now := time.Now()
	b := &strings.Builder{}
	for i, r := range rows {
		fmt.Fprintf(b, "- %s (%s): ", r.c.Host, r.c.Login)
		if r.last == nil {
			b.WriteString("no action since byre started")
		} else {
			fmt.Fprintf(b, "last action %s ago", now.Sub(*r.last).Round(time.Second))
		}
		fmt.Fprintf(b, ", connected %s ago, app built from %s", now.Sub(r.c.ConnectedAt).Round(time.Second), r.c.Commit)
		if i == 0 && r.last != nil {
			b.WriteString(" (most recently active)")
		}
		b.WriteString("\n")
	}
	return b.String(), nil
}

type clientBashArgs struct {
	Client    string `json:"client"`
	Command   string `json:"command"`
	TimeoutMs *int   `json:"timeout_ms"`
}

func (m *Manager) clientBash(ctx context.Context, raw json.RawMessage) (string, error) {
	args := clientBashArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed client_bash call: %w", err)
	}
	if strings.TrimSpace(args.Command) == "" {
		return "", errors.New("command is empty")
	}
	timeout := defaultBashTimeout
	if args.TimeoutMs != nil {
		timeout = time.Duration(*args.TimeoutMs) * time.Millisecond
	}
	if timeout <= 0 || timeout > maxBashTimeout {
		return "", fmt.Errorf("timeout_ms must be between 1 and %d", maxBashTimeout.Milliseconds())
	}

	ms := int(timeout.Milliseconds())
	cctx, cancel := context.WithTimeout(ctx, timeout+replyGrace)
	defer cancel()
	reply, err := m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Bash, Command: &args.Command, TimeoutMs: &ms})
	if err != nil {
		return "", err
	}
	if reply.ExitCode == nil || reply.Output == nil {
		return "", fmt.Errorf("%s's reply has no exit_code or output", args.Client)
	}
	status := fmt.Sprintf("exit status %d", *reply.ExitCode)
	if *reply.ExitCode == -1 {
		status = fmt.Sprintf("killed after %s", timeout)
	}
	return status + "\n" + *reply.Output, nil
}

type clientTransferArgs struct {
	Client string `json:"client"`
	Path   string `json:"path"`
	Dest   string `json:"dest"`
}

func clientPath(name, p string) error {
	if !filepath.IsAbs(p) && !strings.HasPrefix(p, "~/") {
		return fmt.Errorf("%s %q must be absolute or start with ~/", name, p)
	}
	return nil
}

func (m *Manager) clientUpload(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := clientTransferArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed client_upload call: %w", err)
	}
	err = clientPath("dest", args.Dest)
	if err != nil {
		return "", err
	}

	clean := filepath.Clean(args.Path)
	if filepath.IsAbs(clean) {
		return "", fmt.Errorf("path %q must be relative to your worktree", args.Path)
	}
	root, err := os.OpenRoot(l.dir)
	if err != nil {
		return "", err
	}
	_, err = root.Lstat(clean)
	root.Close()
	if err != nil {
		return "", fmt.Errorf("path %q: %w", args.Path, err)
	}

	full := filepath.Join(l.dir, clean)
	id := m.clients.AddTransfer(clients.Transfer{Dir: filepath.Dir(full), Name: filepath.Base(full)})
	defer m.clients.RemoveTransfer(id)

	cctx, cancel := context.WithTimeout(ctx, copyTimeout)
	defer cancel()
	_, err = m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Upload, Transfer: &id, Dest: &args.Dest})
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Copied %s to %s on %s.", clean, args.Dest, args.Client), nil
}

func (m *Manager) clientDownload(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := clientTransferArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed client_download call: %w", err)
	}
	err = clientPath("path", args.Path)
	if err != nil {
		return "", err
	}
	dest := filepath.Clean(args.Dest)
	if filepath.IsAbs(dest) || dest == "." {
		return "", fmt.Errorf("dest %q must be a path inside your worktree, relative to it", args.Dest)
	}

	root, err := os.OpenRoot(l.dir)
	if err != nil {
		return "", err
	}
	defer root.Close()
	parent := filepath.Dir(dest)
	err = root.MkdirAll(parent, 0o755)
	if err != nil {
		return "", fmt.Errorf("dest %q: %w", args.Dest, err)
	}

	tmp, err := os.MkdirTemp("", "byre-download-")
	if err != nil {
		return "", err
	}
	defer os.RemoveAll(tmp)
	id := m.clients.AddTransfer(clients.Transfer{Dir: tmp, Name: "archive.tgz", Receive: true})
	defer m.clients.RemoveTransfer(id)

	cctx, cancel := context.WithTimeout(ctx, copyTimeout)
	defer cancel()
	_, err = m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Download, Path: &args.Path, Transfer: &id})
	if err != nil {
		return "", err
	}

	stage := filepath.Join(parent, ".byre-download-"+id)
	err = root.Mkdir(stage, 0o700)
	if err != nil {
		return "", err
	}
	defer root.RemoveAll(stage)
	out, err := exec.CommandContext(ctx, "tar", "-xzf", filepath.Join(tmp, "archive.tgz"), "-C", filepath.Join(l.dir, stage)).CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("unpack: %w: %s", err, out)
	}
	entries, err := fs.ReadDir(root.FS(), stage)
	if err != nil {
		return "", err
	}
	if len(entries) != 1 {
		return "", fmt.Errorf("unpack: the archive holds %d entries, not 1", len(entries))
	}
	err = root.RemoveAll(dest)
	if err != nil {
		return "", err
	}
	err = root.Rename(filepath.Join(stage, entries[0].Name()), dest)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Copied %s on %s to %s.", args.Path, args.Client, dest), nil
}

func (m *Manager) readClientFile(ctx context.Context, client, path string) ([]byte, error) {
	err := clientPath("path", path)
	if err != nil {
		return nil, err
	}
	cctx, cancel := context.WithTimeout(ctx, fileTimeout)
	defer cancel()
	reply, err := m.clients.Call(cctx, client, api.ClientRequest{Op: api.ReadFile, Path: &path})
	if err != nil {
		return nil, err
	}
	if reply.Data == nil {
		return nil, fmt.Errorf("%s's reply has no data", client)
	}
	return *reply.Data, nil
}

func (m *Manager) writeClientFile(ctx context.Context, client, path string, data []byte) error {
	err := clientPath("path", path)
	if err != nil {
		return err
	}
	cctx, cancel := context.WithTimeout(ctx, fileTimeout)
	defer cancel()
	_, err = m.clients.Call(cctx, client, api.ClientRequest{Op: api.WriteFile, Path: &path, Data: &data})
	return err
}

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

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

func (m *Manager) clientRead(ctx context.Context, raw json.RawMessage) mcp.Result {
	args := clientReadArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return mcp.TextResult(fmt.Sprintf("malformed client_read call: %s", err), true)
	}
	data, err := m.readClientFile(ctx, args.Client, args.Path)
	if err != nil {
		return mcp.TextResult(err.Error(), true)
	}
	mime, ok := imageTypes[strings.ToLower(filepath.Ext(args.Path))]
	if ok {
		return mcp.Result{Content: []mcp.Content{{Type: "image", Data: base64.StdEncoding.EncodeToString(data), MimeType: mime}}}
	}
	text, err := numberLines(args.Path, data, args.Offset, args.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; use client_download to copy it into your worktree", 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
}

type clientWriteArgs struct {
	Client  string `json:"client"`
	Path    string `json:"path"`
	Content string `json:"content"`
}

func (m *Manager) clientWrite(ctx context.Context, raw json.RawMessage) (string, error) {
	args := clientWriteArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed client_write call: %w", err)
	}
	err = m.writeClientFile(ctx, args.Client, args.Path, []byte(args.Content))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Wrote %d bytes to %s on %s.", len(args.Content), args.Path, args.Client), nil
}

type clientEditArgs struct {
	Client    string `json:"client"`
	Path      string `json:"path"`
	OldString string `json:"old_string"`
	NewString string `json:"new_string"`
}

func (m *Manager) clientEdit(ctx context.Context, raw json.RawMessage) (string, error) {
	args := clientEditArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed client_edit call: %w", err)
	}
	if args.OldString == "" {
		return "", errors.New("old_string is empty")
	}
	if args.OldString == args.NewString {
		return "", errors.New("old_string and new_string are the same")
	}
	data, err := m.readClientFile(ctx, args.Client, args.Path)
	if err != nil {
		return "", err
	}
	if !utf8.Valid(data) {
		return "", fmt.Errorf("%s isn't text", args.Path)
	}
	text := string(data)
	n := strings.Count(text, args.OldString)
	if n == 0 {
		return "", fmt.Errorf("old_string isn't in %s", args.Path)
	}
	if n > 1 {
		return "", fmt.Errorf("old_string appears %d times in %s; include more of the text around it so it appears once", n, args.Path)
	}
	err = m.writeClientFile(ctx, args.Client, args.Path, []byte(strings.Replace(text, args.OldString, args.NewString, 1)))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Edited %s on %s.", args.Path, args.Client), nil
}

type clientHandsArgs struct {
	Client    string          `json:"client"`
	Tool      string          `json:"tool"`
	Arguments *map[string]any `json:"arguments"`
}

func (m *Manager) clientHands(ctx context.Context, raw json.RawMessage) mcp.Result {
	args := clientHandsArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return mcp.TextResult(fmt.Sprintf("malformed client_hands call: %s", err), true)
	}
	if args.Arguments == nil {
		args.Arguments = &map[string]any{}
	}

	cctx, cancel := context.WithTimeout(ctx, handsTimeout)
	defer cancel()
	reply, err := m.clients.Call(cctx, args.Client, api.ClientRequest{Op: api.Hands, Tool: &args.Tool, Arguments: args.Arguments})
	if err != nil {
		return mcp.TextResult(err.Error(), true)
	}
	if reply.Result == nil {
		return mcp.TextResult(fmt.Sprintf("%s's reply has no result", args.Client), true)
	}
	res := mcp.Result{}
	err = remarshal(*reply.Result, &res)
	if err != nil {
		return mcp.TextResult(fmt.Sprintf("%s's result isn't an MCP tool result: %s", args.Client, err), true)
	}
	return res
}