byre/internal/agents/preview.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

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

const maxPreviewName = 40

func portEnv(t api.Task) []string {
	return []string{"PORT=" + strconv.Itoa(t.PortBase)}
}

type exposePortArgs struct {
	Name string `json:"name"`
	Port int    `json:"port"`
}

func (m *Manager) exposePortTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "expose_port",
		Description: "Show the human a server you are running, such as a dev server, as a link above the task's tabs that opens it in a new window. " +
			"Start the server yourself first (in the background), listening on 127.0.0.1 (not 0.0.0.0 or localhost) on a port in this task's block. " +
			"byre then serves it over HTTPS on the same port of byre's host through tailscale serve, and returns the URL. " +
			"name labels it in the tab (\"web\", \"storybook\"); exposing a name or port again replaces the earlier one. " +
			"The server stops when your process does (a byre restart, compaction, the task closing); the human sees it as unreachable until you start it again.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"name": {"type": "string", "description": "Short label, shown as the link's text."},
				"port": {"type": "integer", "description": "The port the server listens on, at 127.0.0.1, within this task's block."}
			},
			"required": ["name", "port"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.exposePort(ctx, l, raw)
		},
	}
}

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

	name := strings.TrimSpace(args.Name)
	if name == "" || strings.ContainsAny(name, "\r\n") || utf8.RuneCountInString(name) > maxPreviewName {
		return "", fmt.Errorf("name must be one non-empty line of at most %d characters", maxPreviewName)
	}

	task, err := m.store.GetTask(ctx, l.taskID)
	if err != nil {
		return "", err
	}
	last := task.PortBase + store.PortBlock - 1
	if args.Port < task.PortBase || args.Port > last {
		return "", fmt.Errorf("port %d is outside this task's block, %d to %d", args.Port, task.PortBase, last)
	}

	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(args.Port))
	conn, err := net.DialTimeout("tcp", addr, time.Second)
	if err != nil {
		return "", fmt.Errorf("nothing is listening on %s; start the server listening on 127.0.0.1 first: %w", addr, err)
	}
	conn.Close()

	err = tailscale.Serve(ctx, args.Port)
	if err != nil {
		return "", err
	}

	dropped, err := m.store.SetPreview(ctx, api.Preview{TaskId: task.Id, Name: name, Port: args.Port, CreatedAt: time.Now().UTC()}, agentActor)
	if errors.Is(err, store.ErrNotFound) {
		return "", errors.New("this task is no longer active")
	}
	if err != nil {
		return "", err
	}
	for _, port := range dropped {
		err = tailscale.Off(ctx, port)
		if err != nil {
			slog.Error("stop serving replaced preview", "task", task.Id, "port", port, "err", err)
		}
	}

	return fmt.Sprintf("Exposed %q: https://<byre's host>:%d/ proxies to %s, and the human sees a link to it above the task's tabs.", name, args.Port, addr), nil
}

func (m *Manager) Unexpose(ctx context.Context, taskID string) {
	previews, err := m.store.ListPreviews(ctx, taskID)
	if err != nil {
		slog.Error("list previews", "task", taskID, "err", err)
		return
	}
	for _, p := range previews {
		err = tailscale.Off(ctx, p.Port)
		if err != nil {
			slog.Error("stop serving preview", "task", taskID, "port", p.Port, "err", err)
			continue
		}
		slog.Info("stopped serving preview", "task", taskID, "name", p.Name, "port", p.Port)
	}
}