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/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/clients"
	"g.fc.run/ai/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 id of that running app (pass it as client instead when two apps are open on one Mac), 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, or its id when several are connected from that Mac."},
					"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, or its id when several are connected from that Mac."},
					"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, or its id when several are connected from that Mac."},
					"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, or its id when several are connected from that Mac."},
					"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, or its id when several are connected from that Mac."},
					"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, or its id when several are connected from that Mac."},
					"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, chrome) 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, displays and Chrome tabs are that Mac's, so start with screen_list (or chrome tabs) 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. " +
				"chrome needs Google Chrome installed on that Mac; its trace files are saved there, so fetch one with client_download.",
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {
					"client": {"type": "string", "description": "The client's name from list_clients, or its id when several are connected from that Mac."},
					"tool": {"type": "string", "enum": ["screen_list", "capture", "ax_tree", "act", "browser", "chrome"]},
					"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
	}
	rows := status.Clients
	if len(rows) == 0 {
		return "No client Macs are connected.", nil
	}
	sort.SliceStable(rows, func(i, j int) bool {
		a, b := rows[i].LastActionAt, rows[j].LastActionAt
		return a != nil && (b == nil || a.After(*b))
	})

	now := time.Now()
	b := &strings.Builder{}
	for i, r := range rows {
		fmt.Fprintf(b, "- %s (%s, id %s): ", r.Host, r.Login, r.Id)
		if r.LastActionAt == nil {
			b.WriteString("no action since byre started")
		} else {
			fmt.Fprintf(b, "last action %s ago", now.Sub(*r.LastActionAt).Round(time.Second))
		}
		fmt.Fprintf(b, ", connected %s ago, app built from %s", now.Sub(r.ConnectedAt).Round(time.Second), r.Commit)
		if i == 0 && r.LastActionAt != 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
}