byre/internal/agents/terminal.go

package agents

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/terminal"
)

const (
	typeQuiet   = 500 * time.Millisecond
	typeMaxWait = 10 * time.Second
	typeKeep    = 256 << 10
)

var agentTerminalSize = api.TerminalSize{Cols: 120, Rows: 40}

func (m *Manager) terminalReadTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "terminal_read",
		Description: "Read your task's terminal: the shell in this task's worktree that the human uses from the Terminal tab, and that terminal_type types into. " +
			"Returns its recent output as plain text, escape sequences removed, at most the last 64 KiB. " +
			"Full-screen programs (vim, top, less) come out garbled, since this is the output stream, not a picture of the screen.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {}}`),
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			term := m.terminals.Get(l.taskID)
			if term == nil {
				return "No shell is running in this task's terminal.", nil
			}
			out := terminal.Text(term.Backlog())
			exit, exited := term.ExitStatus()
			if exited {
				out += "\n[shell ended: " + exit + "]"
			}
			return out + inputLine(term), nil
		},
	}
}

func inputLine(term *terminal.Terminal) string {
	line, ok := term.InputLine()
	if !ok {
		return ""
	}
	if line == "" {
		return "\n[input line: empty]"
	}
	return fmt.Sprintf("\n[input line: %q]", line)
}

var terminalKeys = map[string]string{
	"enter":     "\r",
	"escape":    "\x1b",
	"tab":       "\t",
	"backspace": "\x7f",
	"up":        "\x1b[A",
	"down":      "\x1b[B",
	"right":     "\x1b[C",
	"left":      "\x1b[D",
}

func init() {
	for c := 'a'; c <= 'z'; c++ {
		terminalKeys["ctrl-"+string(c)] = string(c - 'a' + 1)
	}
}

type terminalTypeArgs struct {
	Keys   []string `json:"keys"`
	Text   string   `json:"text"`
	Append bool     `json:"append"`
}

func (m *Manager) terminalTypeTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "terminal_type",
		Description: "Type into your task's terminal: the shell in this task's worktree that the human uses from the Terminal tab, starting it when none is running and waiting for its prompt before typing. " +
			"keys are named keys sent first (ctrl-c cancels, ctrl-u clears the input line, enter runs it), then text goes to the shell exactly as keystrokes, so end a command with a real line break to run it (an escape you write in text, such as \"\\u0003\", can arrive as those literal characters, so send control keys through keys), or press enter in a later call. " +
			"When the input line still holds unsent text after keys, text is refused, so that it isn't joined onto someone else's command: clear it with ctrl-u first, or set append to add to it on purpose. " +
			"Replies from this tool and terminal_read end with the shell's current input line when zsh is at its prompt. " +
			"The human sees everything you type and its output. " +
			"To hand the human a command to run (one needing their password, say), type it without the newline and tell them to press Enter in the Terminal tab. In a zsh terminal, the reply then says byre will message you when the shell next returns to its prompt, with the exit status and output, starting a turn if you are idle, so end your turn rather than asking whether they ran it. " +
			"Returns the output that followed, as terminal_read does, once the terminal has been quiet for half a second or after 10 seconds; read the rest of a longer command's output with terminal_read.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"keys": {"type": "array", "items": {"type": "string", "enum": ["ctrl-a", "ctrl-b", "ctrl-c", "ctrl-d", "ctrl-e", "ctrl-f", "ctrl-g", "ctrl-h", "ctrl-i", "ctrl-j", "ctrl-k", "ctrl-l", "ctrl-m", "ctrl-n", "ctrl-o", "ctrl-p", "ctrl-q", "ctrl-r", "ctrl-s", "ctrl-t", "ctrl-u", "ctrl-v", "ctrl-w", "ctrl-x", "ctrl-y", "ctrl-z", "enter", "escape", "tab", "backspace", "up", "down", "left", "right"]}, "description": "Named keys to press, in order, before text."},
				"text": {"type": "string", "description": "The keystrokes to send after keys."},
				"append": {"type": "boolean", "description": "Type text even though the input line already holds unsent text, adding to it."}
			}
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := terminalTypeArgs{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed terminal_type call: %w", err)
			}
			task, err := m.store.GetTask(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			if task.Status != api.TaskStatusActive {
				return "", fmt.Errorf("task %s is %s", task.Id, task.Status)
			}
			term, err := m.terminals.Attach(task, agentTerminalSize)
			if err != nil {
				return "", err
			}
			keys := ""
			for _, k := range args.Keys {
				b, ok := terminalKeys[k]
				if !ok {
					return "", fmt.Errorf("unknown key %q", k)
				}
				keys += b
			}
			term.CancelWake()
			startLine, _ := term.InputLine()
			sent := keys
			out := ""
			if keys != "" || args.Text == "" {
				out, err = typeAndWait(ctx, term, keys)
				if err != nil {
					return "", err
				}
			}
			if args.Text != "" {
				if line, _ := term.InputLine(); line != "" && !args.Append {
					out += fmt.Sprintf("\nThe terminal's input line already has unsent text: %q. Your text was not typed. Send keys [\"ctrl-u\"] to clear it (keys [\"ctrl-c\"] when it continues an unfinished command from earlier lines), or set append to add to it.", line)
				} else {
					more, err := typeAndWait(ctx, term, args.Text)
					if err != nil {
						return "", err
					}
					out += more
					sent += args.Text
				}
			}
			out += inputLine(term)
			submitted := strings.HasSuffix(sent, "\n") || strings.HasSuffix(sent, "\r")
			if !term.Prompts {
				if args.Text != "" {
					out += "\n[this shell isn't zsh, so byre can't see its input line: your text may have joined unsent text already there]"
				}
				if args.Text != "" && !submitted {
					out += "\n[this shell isn't zsh, so byre can't tell when the line is run: ask the human to tell you once they have run it]"
				}
				return out, nil
			}
			if term.WakeOnPrompt(m.wakeOnPrompt(l), !submitted && (startLine != "" || args.Text != "")) {
				out += "\n[byre messages you when the shell is next back at its prompt]"
			}
			return out, nil
		},
	}
}

func (m *Manager) wakeOnPrompt(l *live) func(string, []byte) {
	return func(exit string, out []byte) {
		text := fmt.Sprintf("The line you left unsent in the terminal was run, and the shell is back at its prompt with exit status %s (130 means it was canceled with Ctrl-C). The terminal's output since you typed it:\n\n%s", exit, terminal.Text(out))
		_, _, err := m.Message(context.Background(), l.id, "byre", text, "", "")
		if err != nil {
			slog.Error("wake agent on terminal prompt", "session", l.id, "task", l.taskID, "err", err)
			return
		}
		slog.Info("woke agent on terminal prompt", "session", l.id, "task", l.taskID, "exit", exit)
	}
}

func typeAndWait(ctx context.Context, term *terminal.Terminal, text string) (string, error) {
	backlog, out := term.Subscribe()
	if out == nil {
		exit, _ := term.ExitStatus()
		return "", fmt.Errorf("the shell ended (%s); type again to start a new one", exit)
	}
	defer term.Unsubscribe(out)

	got := []byte{}
	if promptWait(term) > 0 {
		var err error
		got, err = awaitPrompt(ctx, term, out, keepTail(backlog))
		if err != nil {
			return "", err
		}
	}

	stop := atomic.Bool{}
	defer stop.Store(true)
	typed := atomic.Int64{}
	written := make(chan error, 1)
	go func() {
		// One byte per write: a program that stops reading blocks a write, and the PTY can't be given a deadline, so this keeps what arrives after we give up to one byte.
		for i := 0; i < len(text) && !stop.Load(); i++ {
			err := term.Write([]byte{text[i]})
			if err != nil {
				written <- err
				return
			}
			typed.Add(1)
		}
		written <- nil
	}()

	result := func() string {
		res := terminal.Text(got)
		exit, exited := term.ExitStatus()
		if exited {
			res += "\n[shell ended: " + exit + "]"
		}
		if n := typed.Load(); n < int64(len(text)) {
			res += fmt.Sprintf("\n[only %d of %d bytes were typed: the terminal stopped reading input]", n, len(text))
		}
		return res
	}
	deadline := time.After(typeMaxWait)
	quiet := time.NewTimer(typeQuiet)
	defer quiet.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case err := <-written:
			if err != nil {
				return "", err
			}
			written = nil
		case <-deadline:
			return result(), nil
		case <-quiet.C:
			return result(), nil
		case b, open := <-out:
			if !open {
				return result(), nil
			}
			got = keepTail(append(got, b...))
			quiet.Reset(typeQuiet)
		}
	}
}

func keepTail(b []byte) []byte {
	if len(b) > 2*typeKeep {
		return bytes.Clone(b[len(b)-typeKeep:])
	}
	return b
}

func promptWait(term *terminal.Terminal) time.Duration {
	started, last := term.Activity()
	wait := time.Until(started.Add(typeMaxWait))
	if !last.IsZero() {
		wait = min(wait, time.Until(last.Add(typeQuiet)))
	}
	return wait
}

func awaitPrompt(ctx context.Context, term *terminal.Terminal, out chan []byte, got []byte) ([]byte, error) {
	for {
		wait := promptWait(term)
		if wait <= 0 {
			return got, nil
		}
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		case b, open := <-out:
			timer.Stop()
			if !open {
				return got, nil
			}
			got = keepTail(append(got, b...))
		}
	}
}