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/store"
	"g.fc.run/ai/byre/internal/terminal"
)

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

var newTerminalSize = api.TerminalSize{Cols: 80, Rows: 24}

func (s *Server) terminalsChanged(taskID string) {
	s.terminalsMu.Lock()
	defer s.terminalsMu.Unlock()
	list := api.TerminalList{Terminals: []api.TaskTerminal{}}
	for _, term := range s.terminals.List(taskID) {
		list.Terminals = append(list.Terminals, taskTerminal(taskID, term))
	}
	s.store.Broadcast(api.Event{Kind: "terminals.changed", TaskId: &taskID, Actor: "byre", Payload: store.ToPayload(list)})
}

func taskTerminal(taskID string, term *terminal.Terminal) api.TaskTerminal {
	return api.TaskTerminal{Id: term.ID, TaskId: taskID, Name: term.Name, StartedAt: term.Started().UTC()}
}

func (s *Server) ListTerminals(w http.ResponseWriter, r *http.Request) {
	list := api.TerminalList{Terminals: []api.TaskTerminal{}}
	for taskID, terms := range s.terminals.All() {
		for _, term := range terms {
			list.Terminals = append(list.Terminals, taskTerminal(taskID, term))
		}
	}
	writeJSON(w, http.StatusOK, list)
}

func (s *Server) StartTaskTerminal(w http.ResponseWriter, r *http.Request, taskID string) {
	t, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}
	req := api.TerminalStart{}
	if !readJSON(w, r, &req) {
		return
	}
	name, command := "Shell", ""
	if req.CommandId != nil {
		c, err := s.store.SavedCommand(r.Context(), actor(r), *req.CommandId)
		if err != nil {
			writeStoreError(w, err)
			return
		}
		name, command = c.Name, c.Command
	}
	term, err := s.terminals.Start(t, name, command, newTerminalSize)
	if err != nil {
		slog.Error("start terminal", "task", taskID, "err", err)
		writeError(w, http.StatusInternalServerError, err)
		return
	}
	writeJSON(w, http.StatusCreated, taskTerminal(taskID, term))
}

func (s *Server) CloseTaskTerminal(w http.ResponseWriter, r *http.Request, taskID, terminalID string) {
	if !s.terminals.Kill(taskID, terminalID) {
		writeError(w, http.StatusNotFound, fmt.Errorf("no terminal %s in task %s", terminalID, taskID))
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID, terminalID string, params api.GetTaskTerminalParams) {
	_, 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
	}
	term := s.terminals.Get(taskID, terminalID)
	if term == nil {
		writeError(w, http.StatusNotFound, fmt.Errorf("no terminal %s in task %s", terminalID, taskID))
		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.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)
	}
}