byre/internal/agents/preview.go

package agents

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

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

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 on the same port of byre's tailnet node, and the human's app opens it through a local link. " +
			"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 = m.servePreview(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) {
		m.stopPreview(args.Port)
		return "", errors.New("this task is no longer active")
	}
	if err != nil {
		return "", err
	}
	for _, port := range dropped {
		m.stopPreview(port)
	}

	return fmt.Sprintf("Exposed %q: port %d of byre's tailnet node proxies to %s, and the human sees a link to it above the task's tabs.", name, args.Port, addr), nil
}

func (m *Manager) servePreview(port int) error {
	m.stopPreview(port)
	ln, err := m.tailnet.Listen("tcp", ":"+strconv.Itoa(port))
	if err != nil {
		return fmt.Errorf("listen on the tailnet on port %d: %w", port, err)
	}
	lc, err := m.tailnet.LocalClient()
	if err != nil {
		ln.Close()
		return err
	}
	tracked := &previewListener{Listener: ln, conns: map[net.Conn]bool{}}
	m.previewMu.Lock()
	m.previewListeners[port] = tracked
	m.previewMu.Unlock()

	target := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
	proxy := &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.Out.URL.Scheme = "http"
			pr.Out.URL.Host = target
		},
		FlushInterval: -1,
	}
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		who, err := lc.WhoIs(r.Context(), r.RemoteAddr)
		if err != nil {
			http.Error(w, err.Error(), http.StatusForbidden)
			return
		}
		r.Header.Set("Tailscale-User-Login", who.UserProfile.LoginName)
		r.Header.Set("Tailscale-User-Name", who.UserProfile.DisplayName)
		proxy.ServeHTTP(w, r)
	})
	go func() {
		err := (&http.Server{Handler: handler}).Serve(tracked)
		if !errors.Is(err, net.ErrClosed) {
			slog.Error("serve preview", "port", port, "err", err)
		}
	}()
	return nil
}

func (m *Manager) stopPreview(port int) {
	m.previewMu.Lock()
	ln := m.previewListeners[port]
	delete(m.previewListeners, port)
	m.previewMu.Unlock()
	if ln != nil {
		ln.closeAll()
	}
}

type previewListener struct {
	net.Listener
	mu     sync.Mutex
	conns  map[net.Conn]bool
	closed bool
}

func (l *previewListener) Accept() (net.Conn, error) {
	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.closed {
		c.Close()
		return nil, net.ErrClosed
	}
	tc := &previewConn{Conn: c, l: l}
	l.conns[tc] = true
	return tc, nil
}

func (l *previewListener) closeAll() {
	l.Listener.Close()
	l.mu.Lock()
	l.closed = true
	conns := l.conns
	l.conns = map[net.Conn]bool{}
	l.mu.Unlock()
	for c := range conns {
		c.Close()
	}
}

type previewConn struct {
	net.Conn
	l *previewListener
}

func (c *previewConn) Close() error {
	c.l.mu.Lock()
	delete(c.l.conns, c)
	c.l.mu.Unlock()
	return c.Conn.Close()
}

func (m *Manager) recoverPreviews(ctx context.Context) error {
	previews, err := m.store.ActivePreviews(ctx)
	if err != nil {
		return err
	}
	for _, p := range previews {
		err = m.servePreview(p.Port)
		if err != nil {
			slog.Error("serve preview", "task", p.TaskId, "port", p.Port, "err", err)
		}
	}
	return 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 {
		m.stopPreview(p.Port)
		slog.Info("stopped serving preview", "task", taskID, "name", p.Name, "port", p.Port)
	}
}