byre/internal/server/terminal.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"

	"github.com/coder/websocket"

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

const (
	terminalReadLimit = 1 << 20
	terminalExited    = websocket.StatusCode(4000)
)

func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID string, params api.GetTaskTerminalParams) {
	t, _, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}
	size := api.TerminalSize{Cols: params.Cols, Rows: params.Rows}
	if !terminal.ValidSize(size) {
		writeError(w, http.StatusBadRequest, fmt.Errorf("cols and rows must be between 1 and %d", terminal.MaxSize))
		return
	}

	conn, err := websocket.Accept(w, r, nil)
	if err != nil {
		slog.Warn("accept terminal websocket", "task", taskID, "err", err)
		return
	}
	defer conn.CloseNow()
	conn.SetReadLimit(terminalReadLimit)

	term, err := s.terminals.Attach(t, size)
	if err != nil {
		slog.Error("start terminal", "task", taskID, "err", err)
		conn.Close(websocket.StatusInternalError, "couldn't start the shell; see byre's log")
		return
	}
	term.Resize(size)

	backlog, out := term.Subscribe()
	if out == nil {
		exit, _ := term.ExitStatus()
		conn.Close(terminalExited, exit)
		return
	}
	defer term.Unsubscribe(out)

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	go func() {
		defer cancel()
		terminalInput(ctx, conn, term, func() { s.acted(actor(r), r) })
	}()

	if len(backlog) > 0 {
		err = conn.Write(ctx, websocket.MessageBinary, backlog)
		if err != nil {
			return
		}
	}
	for {
		select {
		case <-ctx.Done():
			return
		case b, open := <-out:
			if open {
				err = conn.Write(ctx, websocket.MessageBinary, b)
				if err != nil {
					return
				}
				continue
			}
			exit, exited := term.ExitStatus()
			if exited {
				conn.Close(terminalExited, exit)
				return
			}
			conn.Close(websocket.StatusTryAgainLater, "fell behind the terminal's output")
			return
		}
	}
}

func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal.Terminal, acted func()) {
	for {
		typ, b, err := conn.Read(ctx)
		if err != nil {
			return
		}
		if typ == websocket.MessageBinary {
			acted()
			err = term.Write(b)
			if err != nil {
				slog.Debug("terminal input", "err", err)
			}
			continue
		}
		size := api.TerminalSize{}
		err = json.Unmarshal(b, &size)
		if err != nil || !terminal.ValidSize(size) {
			conn.Close(websocket.StatusUnsupportedData, "resize must be a TerminalSize")
			return
		}
		term.Resize(size)
	}
}