Let task agents use client Macs directly: bash, copy and hands

internal/clients/clients.go

@@ -0,0 +1,202 @@
+package clients
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"slices"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/coder/websocket"
+
+	"g.fc.run/misc/byre/internal/api"
+)
+
+const (
+	readLimit    = 256 << 20
+	pingInterval = 30 * time.Second
+	pingTimeout  = 15 * time.Second
+)
+
+type Client struct {
+	Login       string
+	Host        string
+	Commit      string
+	ConnectedAt time.Time
+}
+
+type Transfer struct {
+	Dir  string
+	Name string
+}
+
+type conn struct {
+	Client
+	ws      *websocket.Conn
+	done    chan struct{}
+	mu      sync.Mutex
+	pending map[string]chan api.ClientReply
+}
+
+type Registry struct {
+	mu        sync.Mutex
+	conns     map[string]*conn
+	transfers map[string]Transfer
+}
+
+func New() *Registry {
+	return &Registry{conns: map[string]*conn{}, transfers: map[string]Transfer{}}
+}
+
+func (r *Registry) Serve(ctx context.Context, ws *websocket.Conn, client Client) {
+	ws.SetReadLimit(readLimit)
+	c := &conn{Client: client, ws: ws, done: make(chan struct{}), pending: map[string]chan api.ClientReply{}}
+
+	r.mu.Lock()
+	old := r.conns[client.Host]
+	r.conns[client.Host] = c
+	r.mu.Unlock()
+	if old != nil {
+		old.ws.Close(websocket.StatusGoingAway, "replaced by a newer connection from this host")
+	}
+	slog.Info("client connected", "host", client.Host, "login", client.Login, "commit", client.Commit)
+
+	defer func() {
+		r.mu.Lock()
+		if r.conns[client.Host] == c {
+			delete(r.conns, client.Host)
+		}
+		r.mu.Unlock()
+		close(c.done)
+		slog.Info("client disconnected", "host", client.Host, "login", client.Login)
+	}()
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	go c.ping(ctx)
+
+	for {
+		_, b, err := ws.Read(ctx)
+		if err != nil {
+			return
+		}
+		reply := api.ClientReply{}
+		err = json.Unmarshal(b, &reply)
+		if err != nil {
+			ws.Close(websocket.StatusUnsupportedData, "messages must be ClientReply")
+			return
+		}
+		c.mu.Lock()
+		ch := c.pending[reply.Id]
+		delete(c.pending, reply.Id)
+		c.mu.Unlock()
+		if ch == nil {
+			slog.Warn("client reply for no request", "host", client.Host, "id", reply.Id)
+			continue
+		}
+		ch <- reply
+	}
+}
+
+func (c *conn) ping(ctx context.Context) {
+	t := time.NewTicker(pingInterval)
+	defer t.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-t.C:
+			pctx, cancel := context.WithTimeout(ctx, pingTimeout)
+			err := c.ws.Ping(pctx)
+			cancel()
+			if err != nil {
+				c.ws.Close(websocket.StatusGoingAway, "ping timed out")
+				return
+			}
+		}
+	}
+}
+
+func (r *Registry) List() []Client {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	out := []Client{}
+	for _, c := range r.conns {
+		out = append(out, c.Client)
+	}
+	sort.Slice(out, func(i, j int) bool { return out[i].Host < out[j].Host })
+	return out
+}
+
+func (r *Registry) Call(ctx context.Context, host string, req api.ClientRequest) (api.ClientReply, error) {
+	r.mu.Lock()
+	c := r.conns[host]
+	hosts := []string{}
+	for h := range r.conns {
+		hosts = append(hosts, h)
+	}
+	r.mu.Unlock()
+	if c == nil {
+		slices.Sort(hosts)
+		return api.ClientReply{}, fmt.Errorf("no client named %q is connected; connected: %s", host, strings.Join(hosts, ", "))
+	}
+
+	req.Id = rand.Text()
+	ch := make(chan api.ClientReply, 1)
+	c.mu.Lock()
+	c.pending[req.Id] = ch
+	c.mu.Unlock()
+	defer func() {
+		c.mu.Lock()
+		delete(c.pending, req.Id)
+		c.mu.Unlock()
+	}()
+
+	data, err := json.Marshal(req)
+	if err != nil {
+		return api.ClientReply{}, err
+	}
+	err = c.ws.Write(ctx, websocket.MessageText, data)
+	if err != nil {
+		return api.ClientReply{}, fmt.Errorf("send to %s: %w", host, err)
+	}
+
+	select {
+	case reply := <-ch:
+		if reply.Error != nil {
+			return reply, errors.New(*reply.Error)
+		}
+		return reply, nil
+	case <-c.done:
+		return api.ClientReply{}, fmt.Errorf("%s disconnected before replying", host)
+	case <-ctx.Done():
+		return api.ClientReply{}, fmt.Errorf("%s didn't reply: %w", host, ctx.Err())
+	}
+}
+
+func (r *Registry) AddTransfer(t Transfer) string {
+	id := rand.Text()
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.transfers[id] = t
+	return id
+}
+
+func (r *Registry) RemoveTransfer(id string) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	delete(r.transfers, id)
+}
+
+func (r *Registry) Transfer(id string) (Transfer, bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	t, ok := r.transfers[id]
+	return t, ok
+}