iop/apps/node/internal/terminal/session.go

package terminal

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/creack/pty"
)

const (
	terminalInputDelay = 2 * time.Millisecond
)

type Options struct {
	Command string
	Args    []string
	Env     []string
	Rows    uint16
	Cols    uint16
	Dir     string
}

type Output struct {
	Text string
}

type Snapshot struct {
	Tail string
}

type Session interface {
	Output() <-chan Output
	Done() <-chan error
	WritePrompt(ctx context.Context, prompt string) error
	WriteInput(ctx context.Context, data []byte) error
	Resize(rows, cols uint16) error
	Signal(sig os.Signal) error
	Snapshot() Snapshot
	Close() error
}

type sessionImpl struct {
	cmd      *exec.Cmd
	input    io.WriteCloser
	output   io.ReadCloser
	tail     *TailBuffer
	outputCh chan Output
	doneCh   chan error
	closed   bool
	mu       sync.Mutex
}

func StartSession(ctx context.Context, opts Options) (Session, error) {
	cmd := exec.Command(opts.Command, opts.Args...)
	if len(opts.Env) > 0 {
		cmd.Env = append(cmd.Environ(), opts.Env...)
	}
	if opts.Dir != "" {
		cmd.Dir = opts.Dir
	}
	rows := opts.Rows
	if rows == 0 {
		rows = 80
	}
	cols := opts.Cols
	if cols == 0 {
		cols = 240
	}
	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{
		Rows: rows,
		Cols: cols,
	})
	if err != nil {
		return nil, fmt.Errorf("pty start: %w", err)
	}

	s := &sessionImpl{
		cmd:      cmd,
		input:    ptmx,
		output:   ptmx,
		tail:     NewTailBuffer(2048),
		outputCh: make(chan Output, 1024),
		doneCh:   make(chan error, 1),
	}

	go s.readLoop()

	return s, nil
}

func (s *sessionImpl) Output() <-chan Output {
	return s.outputCh
}

func (s *sessionImpl) Done() <-chan error {
	return s.doneCh
}

func (s *sessionImpl) readLoop() {
	buf := make([]byte, 4096)
	for {
		n, err := s.output.Read(buf)
		if n > 0 {
			text := string(buf[:n])
			s.tail.Append(text)
			s.outputCh <- Output{Text: text}
		}
		if err != nil {
			s.doneCh <- s.cmd.Wait()
			close(s.outputCh)
			return
		}
	}
}

func (s *sessionImpl) WritePrompt(ctx context.Context, prompt string) error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return errors.New("terminal session closed")
	}
	s.mu.Unlock()

	for _, r := range prompt {
		s.mu.Lock()
		if s.closed {
			s.mu.Unlock()
			return errors.New("terminal session closed")
		}
		_, err := io.WriteString(s.input, string(r))
		s.mu.Unlock()
		if err != nil {
			return err
		}
		timer := time.NewTimer(terminalInputDelay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("terminal session closed")
	}
	_, err := io.WriteString(s.input, "\r")
	return err
}

func (s *sessionImpl) Snapshot() Snapshot {
	return Snapshot{
		Tail: s.tail.String(),
	}
}

func (s *sessionImpl) Resize(rows, cols uint16) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("terminal session closed")
	}
	if rows == 0 || cols == 0 {
		return errors.New("invalid terminal size: rows and cols must be greater than 0")
	}
	f, ok := s.input.(*os.File)
	if !ok {
		return errors.New("terminal input is not a file")
	}
	return pty.Setsize(f, &pty.Winsize{Rows: rows, Cols: cols})
}

func (s *sessionImpl) WriteInput(ctx context.Context, data []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("terminal session closed")
	}
	_, err := s.input.Write(data)
	return err
}

func (s *sessionImpl) Signal(sig os.Signal) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("terminal session closed")
	}
	if s.cmd == nil || s.cmd.Process == nil {
		return errors.New("no process to signal")
	}
	return s.cmd.Process.Signal(sig)
}

func (s *sessionImpl) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	err := s.input.Close()
	if s.cmd != nil && s.cmd.Process != nil {
		_ = s.cmd.Process.Kill()
	}
	if isAlreadyClosedError(err) {
		return nil
	}
	return err
}

func isAlreadyClosedError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, os.ErrClosed) {
		return true
	}
	errStr := err.Error()
	return strings.Contains(errStr, "file already closed") ||
		strings.Contains(errStr, "use of closed file")
}

// TailBuffer holds a bounded buffer of terminal output.
type TailBuffer struct {
	mu  sync.Mutex
	buf strings.Builder
	max int
}

// NewTailBuffer creates a new TailBuffer with the given maximum capacity in bytes.
func NewTailBuffer(max int) *TailBuffer {
	return &TailBuffer{max: max}
}

// Append appends a string to the buffer, truncating from the front if it exceeds capacity.
func (tb *TailBuffer) Append(s string) {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	tb.buf.WriteString(s)
	raw := tb.buf.String()
	if len(raw) <= tb.max {
		return
	}
	tb.buf.Reset()
	tb.buf.WriteString(raw[len(raw)-tb.max:])
}

// String returns the contents of the buffer.
func (tb *TailBuffer) String() string {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	return tb.buf.String()
}