overlord/Overlord-Client/cmd/agent/runtime/console.go

package runtime

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"sync"
	"syscall"

	"overlord-client/cmd/agent/wire"

	"github.com/creack/pty"
)

type ConsoleHub struct {
	env      *Env
	mu       sync.Mutex
	sessions map[string]*ConsoleSession
}

type ConsoleSession struct {
	id     string
	cmd    *exec.Cmd
	pty    *os.File
	cancel context.CancelFunc
	once   sync.Once
}

type ConsoleStartRequest struct {
	SessionID string
	Cols      int
	Rows      int
}

func NewConsoleHub(env *Env) *ConsoleHub {
	return &ConsoleHub{env: env, sessions: make(map[string]*ConsoleSession)}
}

func (h *ConsoleHub) Start(ctx context.Context, req ConsoleStartRequest) error {
	if req.SessionID == "" {
		return fmt.Errorf("missing session id")
	}

	cols := uint16(req.Cols)
	rows := uint16(req.Rows)
	if cols == 0 {
		cols = 120
	}
	if rows == 0 {
		rows = 36
	}

	h.mu.Lock()
	if existing, ok := h.sessions[req.SessionID]; ok {
		existing.close()
		delete(h.sessions, req.SessionID)
	}
	h.mu.Unlock()

	shell := detectShell()
	sessionCtx, cancel := context.WithCancel(ctx)
	cmd := exec.CommandContext(sessionCtx, shell[0], shell[1:]...)
	cmd.Env = os.Environ()

	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: cols, Rows: rows})
	if err != nil {
		cancel()
		h.emitError(ctx, req.SessionID, err)
		return err
	}

	sess := &ConsoleSession{id: req.SessionID, cmd: cmd, pty: ptmx, cancel: cancel}
	h.mu.Lock()
	h.sessions[req.SessionID] = sess
	h.mu.Unlock()

	go h.forwardOutput(sessionCtx, sess)
	return nil
}

func (h *ConsoleHub) Write(ctx context.Context, sessionID string, data string) error {
	sess := h.get(sessionID)
	if sess == nil {
		return fmt.Errorf("session not found")
	}
	_, err := sess.pty.Write([]byte(data))
	return err
}

func (h *ConsoleHub) Resize(sessionID string, cols, rows int) error {
	sess := h.get(sessionID)
	if sess == nil {
		return fmt.Errorf("session not found")
	}
	if cols <= 0 {
		cols = 120
	}
	if rows <= 0 {
		rows = 36
	}
	return pty.Setsize(sess.pty, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows)})
}

func (h *ConsoleHub) Stop(sessionID string) {
	h.mu.Lock()
	sess := h.sessions[sessionID]
	delete(h.sessions, sessionID)
	h.mu.Unlock()
	if sess != nil {
		sess.close()
	}
}

func (h *ConsoleHub) StopAll() {
	if h == nil {
		return
	}
	h.mu.Lock()
	ids := make([]string, 0, len(h.sessions))
	for id := range h.sessions {
		ids = append(ids, id)
	}
	h.mu.Unlock()
	for _, id := range ids {
		h.Stop(id)
	}
}

func (h *ConsoleHub) forwardOutput(ctx context.Context, sess *ConsoleSession) {
	buf := make([]byte, 4096)
	for {
		n, err := sess.pty.Read(buf)
		if n > 0 {
			chunk := make([]byte, n)
			copy(chunk, buf[:n])
			_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sess.id, Data: chunk})
		}
		if err != nil {
			if err != io.EOF {
				h.emitError(ctx, sess.id, err)
			}
			break
		}
	}

	if sess.cmd != nil {
		_ = sess.cmd.Wait()
	}
	exitCode := extractExitCode(sess.cmd)
	_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sess.id, ExitCode: &exitCode})

	h.mu.Lock()
	delete(h.sessions, sess.id)
	h.mu.Unlock()
	sess.close()
}

func (h *ConsoleHub) emitError(ctx context.Context, sessionID string, err error) {
	_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sessionID, Error: err.Error()})
}

func (h *ConsoleHub) get(sessionID string) *ConsoleSession {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.sessions[sessionID]
}

func (sess *ConsoleSession) close() {
	sess.once.Do(func() {
		if sess.cancel != nil {
			sess.cancel()
		}
		if sess.cmd != nil && sess.cmd.Process != nil {
			_ = sess.cmd.Process.Kill()
		}
		if sess.pty != nil {
			_ = sess.pty.Close()
		}
	})
}

func detectShell() []string {
	if runtime.GOOS == "windows" {
		if c := os.Getenv("COMSPEC"); c != "" {
			return []string{c}
		}
		return []string{"cmd.exe"}
	}
	if sh := os.Getenv("SHELL"); sh != "" {
		return []string{sh, "-l"}
	}
	if runtime.GOOS == "darwin" {
		return []string{"/bin/zsh", "-l"}
	}
	return []string{"/bin/bash", "-l"}
}

func extractExitCode(cmd *exec.Cmd) int {
	if cmd == nil || cmd.ProcessState == nil {
		return -1
	}
	if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
		return status.ExitStatus()
	}
	return -1
}