Let task agents read and type into the task's terminal
This commit is contained in:
Ian Gulliver
+13
-191
@@ -1,207 +1,31 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"syscall"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
"github.com/creack/pty"
|
||||
|
||||
"g.fc.run/misc/byre/internal/api"
|
||||
"g.fc.run/misc/byre/internal/terminal"
|
||||
)
|
||||
|
||||
const (
|
||||
terminalBacklog = 256 << 10
|
||||
terminalClientBuf = 256
|
||||
terminalReadLimit = 1 << 20
|
||||
terminalMaxSize = 1 << 15
|
||||
terminalExited = websocket.StatusCode(4000)
|
||||
)
|
||||
|
||||
type terminals struct {
|
||||
mu sync.Mutex
|
||||
byTask map[string]*terminal
|
||||
}
|
||||
|
||||
type terminal struct {
|
||||
ptmx *os.File
|
||||
cmd *exec.Cmd
|
||||
|
||||
mu sync.Mutex
|
||||
backlog []byte
|
||||
clients map[chan []byte]bool
|
||||
exit string
|
||||
exited bool
|
||||
}
|
||||
|
||||
func newTerminals() *terminals {
|
||||
return &terminals{byTask: map[string]*terminal{}}
|
||||
}
|
||||
|
||||
func (ts *terminals) attach(t api.Task, size api.TerminalSize) (*terminal, error) {
|
||||
ts.mu.Lock()
|
||||
defer ts.mu.Unlock()
|
||||
|
||||
term := ts.byTask[t.Id]
|
||||
if term != nil {
|
||||
return term, nil
|
||||
}
|
||||
|
||||
shell := os.Getenv("SHELL")
|
||||
if shell == "" {
|
||||
return nil, errors.New("SHELL is not set in byre's environment")
|
||||
}
|
||||
cmd := exec.Command(shell, "-l")
|
||||
cmd.Dir = t.WorktreePath
|
||||
cmd.Env = append(os.Environ(), "TERM=xterm-256color", "COLORTERM=truecolor", fmt.Sprintf("PORT=%d", t.PortBase))
|
||||
ptmx, err := pty.StartWithSize(cmd, winsize(size))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("start %s: %w", shell, err)
|
||||
}
|
||||
|
||||
term = &terminal{ptmx: ptmx, cmd: cmd, clients: map[chan []byte]bool{}}
|
||||
ts.byTask[t.Id] = term
|
||||
slog.Info("terminal started", "task", t.Id, "shell", shell, "pid", cmd.Process.Pid)
|
||||
go func() {
|
||||
term.pump()
|
||||
ts.mu.Lock()
|
||||
if ts.byTask[t.Id] == term {
|
||||
delete(ts.byTask, t.Id)
|
||||
}
|
||||
ts.mu.Unlock()
|
||||
slog.Info("terminal ended", "task", t.Id, "exit", term.exit)
|
||||
}()
|
||||
return term, nil
|
||||
}
|
||||
|
||||
func (ts *terminals) close(taskID string) {
|
||||
ts.mu.Lock()
|
||||
term := ts.byTask[taskID]
|
||||
delete(ts.byTask, taskID)
|
||||
ts.mu.Unlock()
|
||||
if term != nil {
|
||||
term.kill()
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *terminals) closeAll() {
|
||||
ts.mu.Lock()
|
||||
all := ts.byTask
|
||||
ts.byTask = map[string]*terminal{}
|
||||
ts.mu.Unlock()
|
||||
for _, term := range all {
|
||||
term.kill()
|
||||
}
|
||||
}
|
||||
|
||||
func winsize(size api.TerminalSize) *pty.Winsize {
|
||||
return &pty.Winsize{Cols: uint16(size.Cols), Rows: uint16(size.Rows)}
|
||||
}
|
||||
|
||||
func validSize(size api.TerminalSize) bool {
|
||||
return size.Cols >= 1 && size.Rows >= 1 && size.Cols <= terminalMaxSize && size.Rows <= terminalMaxSize
|
||||
}
|
||||
|
||||
func (term *terminal) pump() {
|
||||
buf := make([]byte, 32<<10)
|
||||
for {
|
||||
n, err := term.ptmx.Read(buf)
|
||||
if n > 0 {
|
||||
term.output(bytes.Clone(buf[:n]))
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
err := term.cmd.Wait()
|
||||
term.ptmx.Close()
|
||||
exit := fmt.Sprint(err)
|
||||
if term.cmd.ProcessState != nil {
|
||||
exit = term.cmd.ProcessState.String()
|
||||
}
|
||||
|
||||
term.mu.Lock()
|
||||
defer term.mu.Unlock()
|
||||
term.exit = exit
|
||||
term.exited = true
|
||||
for c := range term.clients {
|
||||
close(c)
|
||||
}
|
||||
term.clients = nil
|
||||
}
|
||||
|
||||
func (term *terminal) output(b []byte) {
|
||||
term.mu.Lock()
|
||||
defer term.mu.Unlock()
|
||||
|
||||
term.backlog = append(term.backlog, b...)
|
||||
if len(term.backlog) > 2*terminalBacklog {
|
||||
term.backlog = bytes.Clone(term.backlog[len(term.backlog)-terminalBacklog:])
|
||||
}
|
||||
for c := range term.clients {
|
||||
select {
|
||||
case c <- b:
|
||||
default:
|
||||
delete(term.clients, c)
|
||||
close(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (term *terminal) subscribe() ([]byte, chan []byte) {
|
||||
term.mu.Lock()
|
||||
defer term.mu.Unlock()
|
||||
|
||||
if term.exited {
|
||||
return nil, nil
|
||||
}
|
||||
c := make(chan []byte, terminalClientBuf)
|
||||
term.clients[c] = true
|
||||
return bytes.Clone(term.backlog[max(0, len(term.backlog)-terminalBacklog):]), c
|
||||
}
|
||||
|
||||
func (term *terminal) unsubscribe(c chan []byte) {
|
||||
term.mu.Lock()
|
||||
defer term.mu.Unlock()
|
||||
|
||||
if term.clients[c] {
|
||||
delete(term.clients, c)
|
||||
close(c)
|
||||
}
|
||||
}
|
||||
|
||||
func (term *terminal) exitStatus() (string, bool) {
|
||||
term.mu.Lock()
|
||||
defer term.mu.Unlock()
|
||||
return term.exit, term.exited
|
||||
}
|
||||
|
||||
func (term *terminal) kill() {
|
||||
err := syscall.Kill(-term.cmd.Process.Pid, syscall.SIGHUP)
|
||||
if err != nil && !errors.Is(err, syscall.ESRCH) {
|
||||
slog.Warn("hang up terminal", "pid", term.cmd.Process.Pid, "err", err)
|
||||
}
|
||||
term.ptmx.Close()
|
||||
}
|
||||
|
||||
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 !validSize(size) {
|
||||
writeError(w, http.StatusBadRequest, fmt.Errorf("cols and rows must be between 1 and %d", terminalMaxSize))
|
||||
if !terminal.ValidSize(size) {
|
||||
writeError(w, http.StatusBadRequest, fmt.Errorf("cols and rows must be between 1 and %d", terminal.MaxSize))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -213,20 +37,21 @@ func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID
|
||||
defer conn.CloseNow()
|
||||
conn.SetReadLimit(terminalReadLimit)
|
||||
|
||||
term, err := s.terminals.attach(t, size)
|
||||
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()
|
||||
backlog, out := term.Subscribe()
|
||||
if out == nil {
|
||||
exit, _ := term.exitStatus()
|
||||
exit, _ := term.ExitStatus()
|
||||
conn.Close(terminalExited, exit)
|
||||
return
|
||||
}
|
||||
defer term.unsubscribe(out)
|
||||
defer term.Unsubscribe(out)
|
||||
|
||||
ctx, cancel := context.WithCancel(r.Context())
|
||||
defer cancel()
|
||||
@@ -253,7 +78,7 @@ func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID
|
||||
}
|
||||
continue
|
||||
}
|
||||
exit, exited := term.exitStatus()
|
||||
exit, exited := term.ExitStatus()
|
||||
if exited {
|
||||
conn.Close(terminalExited, exit)
|
||||
return
|
||||
@@ -264,7 +89,7 @@ func (s *Server) GetTaskTerminal(w http.ResponseWriter, r *http.Request, taskID
|
||||
}
|
||||
}
|
||||
|
||||
func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal, acted func()) {
|
||||
func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal.Terminal, acted func()) {
|
||||
for {
|
||||
typ, b, err := conn.Read(ctx)
|
||||
if err != nil {
|
||||
@@ -272,7 +97,7 @@ func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal, ac
|
||||
}
|
||||
if typ == websocket.MessageBinary {
|
||||
acted()
|
||||
_, err = term.ptmx.Write(b)
|
||||
err = term.Write(b)
|
||||
if err != nil {
|
||||
slog.Debug("terminal input", "err", err)
|
||||
}
|
||||
@@ -280,13 +105,10 @@ func terminalInput(ctx context.Context, conn *websocket.Conn, term *terminal, ac
|
||||
}
|
||||
size := api.TerminalSize{}
|
||||
err = json.Unmarshal(b, &size)
|
||||
if err != nil || !validSize(size) {
|
||||
if err != nil || !terminal.ValidSize(size) {
|
||||
conn.Close(websocket.StatusUnsupportedData, "resize must be a TerminalSize")
|
||||
return
|
||||
}
|
||||
err = pty.Setsize(term.ptmx, winsize(size))
|
||||
if err != nil {
|
||||
slog.Debug("terminal resize", "err", err)
|
||||
}
|
||||
term.Resize(size)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user