codespace/internal/shell/manager.go

package shell

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
	"syscall"
	"time"

	"codespace/internal/util"

	"github.com/creack/pty"
	"github.com/google/uuid"
)

const (
	outputBufferSize = 4 * 1024
)

// DefaultShellCommand is the default command used to launch a shell.
const DefaultShellCommand = "bash"

// ShellInfo describes a shell instance for listing.
type ShellInfo struct {
	ShellID     string
	WorkspaceID string
	Running     bool
	PID         int
	CreatedAt   time.Time
	ExitCode    int
	Signal      string
}

// Manager manages interactive shells per workspace.
type Manager interface {
	Start(workspaceID string, workspaceRoot string) (string, error)
	Stop(workspaceID string, shellID string) error
	Restart(workspaceID string, shellID string, workspaceRoot string) (string, error)
	Status(workspaceID string, shellID string) Status
	Subscribe(workspaceID string, shellID string) (Subscription, error)
	Stdin(workspaceID string, shellID string) (io.WriteCloser, error)
	ExitStatus(workspaceID string, shellID string) (ExitInfo, error)
	Resize(workspaceID string, shellID string, cols, rows int) error
	List(workspaceID string) []ShellInfo
}

// LocalManager implements Manager using local OS processes.
//
// Concurrency model:
//   - sessionsMu is an RWMutex. Read-only paths (Status, Subscribe lookup,
//     Stdin lookup, ExitStatus lookup, Resize) take RLock; mutating paths
//     (Start, Stop, Restart, List) take Lock.
//   - exited is a sync.Map keyed by shellID (a UUID, globally unique).
//     waitExit is the only goroutine that calls sess.Cmd.Wait() and the only
//     one that writes here. Reads from anywhere are lock-free.
type LocalManager struct {
	command    string
	args       []string
	sessionsMu sync.RWMutex
	sessions   map[string]map[string]*Session // workspaceID -> shellID -> Session
	shellOrder map[string][]string            // workspaceID -> shellIDs in creation order
	exited     sync.Map                       // shellID -> struct{}
}

// NewManager creates a LocalManager with the given command and arguments.
// If command is empty, it defaults to "bash".
func NewManager(command string, args []string) *LocalManager {
	return &LocalManager{
		command:    normalizeCommand(command),
		args:       args,
		sessions:   make(map[string]map[string]*Session),
		shellOrder: make(map[string][]string),
	}
}

func normalizeCommand(command string) string {
	if command == "" {
		return DefaultShellCommand
	}
	return command
}

// IsExited reports whether the shell has exited. Lock-free; safe from any goroutine.
func (m *LocalManager) IsExited(shellID string) bool {
	_, ok := m.exited.Load(shellID)
	return ok
}

// MarkAsExited records the shell as exited. Called only from waitExit.
func (m *LocalManager) MarkAsExited(shellID string) {
	m.exited.Store(shellID, struct{}{})
}

// ClearExited removes the exited marker. Called from Start and Stop/Restart.
func (m *LocalManager) ClearExited(shellID string) {
	m.exited.Delete(shellID)
}

// Start launches a new shell for the given workspace and returns its shellID.
func (m *LocalManager) Start(workspaceID string, workspaceRoot string) (string, error) {
	m.sessionsMu.Lock()
	defer m.sessionsMu.Unlock()

	shellID := uuid.NewString()

	cmd := exec.Command(m.command, m.args...)
	cmd.Dir = workspaceRoot
	cmd.Env = append(os.Environ(),
		"TERM=xterm-256color",
		fmt.Sprintf("HOME=%s", workspaceRoot),
	)

	ptyF, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: 24, Cols: 80, X: 0, Y: 0})
	if err != nil {
		return "", util.Wrap(util.CodeInternal, "failed to start shell on pty", err)
	}

	sess := &Session{
		ShellID:     shellID,
		WorkspaceID: workspaceID,
		Root:        workspaceRoot,
		Cmd:         cmd,
		Stdin:       ptyF,
		Subscribers: make(map[Subscription]struct{}),
		CreatedAt:   time.Now().UTC(),
	}

	if m.sessions[workspaceID] == nil {
		m.sessions[workspaceID] = make(map[string]*Session)
	}
	m.sessions[workspaceID][shellID] = sess

	m.shellOrder[workspaceID] = append(m.shellOrder[workspaceID], shellID)

	m.ClearExited(shellID)

	outputDone := make(chan struct{})
	go m.captureOutput(sess, ptyF, outputDone)
	go m.waitExit(sess, outputDone)

	return shellID, nil
}

// captureOutput reads from the PTY master and fans out each chunk to all
// subscribers using non-blocking sends.
func (m *LocalManager) captureOutput(sess *Session, ptyF *os.File, done chan<- struct{}) {
	defer close(done)
	defer ptyF.Close()

	buf := make([]byte, outputBufferSize)

	for {
		n, err := ptyF.Read(buf)
		if n > 0 {
			chunk := make([]byte, n)
			copy(chunk, buf[:n])

			sess.mu.Lock()
			for sub := range sess.Subscribers {
				sub.(*subscription).send(chunk)
			}
			sess.mu.Unlock()
		}
		if err != nil {
			if err != io.EOF {
				// Ignore read errors; the PTY is closing.
			}
			break
		}
	}
}

// waitExit waits for the shell to exit, then records the exit status and
// closes every subscriber channel exactly once. The PTY master is closed by
// captureOutput once reads finish.
func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) {
	_ = sess.Cmd.Wait()
	<-outputDone

	m.MarkAsExited(sess.ShellID)

	sess.mu.Lock()
	defer sess.mu.Unlock()

	if sess.Cmd.ProcessState != nil {
		sess.Exit.Code = sess.Cmd.ProcessState.ExitCode()
		if ws, ok := sess.Cmd.ProcessState.Sys().(syscall.WaitStatus); ok && ws.Signaled() {
			sess.Exit.Signal = ws.Signal().String()
		}
	}

	for sub := range sess.Subscribers {
		sub.(*subscription).closeChan()
	}
}

// Stop kills the shell for the given workspace and shellID.
// Returns CodeNotFound if no session exists.
func (m *LocalManager) Stop(workspaceID string, shellID string) error {
	m.sessionsMu.Lock()
	defer m.sessionsMu.Unlock()

	ws, ok := m.sessions[workspaceID]
	if !ok {
		return util.New(util.CodeNotFound, "shell not found")
	}
	sess, ok := ws[shellID]
	if !ok {
		return util.New(util.CodeNotFound, "shell not found")
	}

	if sess.Cmd.Process != nil {
		_ = sess.Cmd.Process.Kill()
	}

	delete(ws, shellID)
	if len(ws) == 0 {
		delete(m.sessions, workspaceID)
	}

	m.ClearExited(shellID)
	return nil
}

// Restart stops the identified shell (if running) then starts a new one.
func (m *LocalManager) Restart(workspaceID string, shellID string, workspaceRoot string) (string, error) {
	m.sessionsMu.Lock()
	if ws, ok := m.sessions[workspaceID]; ok {
		if sess, ok := ws[shellID]; ok {
			if sess.Cmd.Process != nil {
				_ = sess.Cmd.Process.Kill()
			}
			delete(ws, shellID)
			if len(ws) == 0 {
				delete(m.sessions, workspaceID)
			}
			m.ClearExited(shellID)
		}
	}
	m.sessionsMu.Unlock()

	return m.Start(workspaceID, workspaceRoot)
}

// Resize resizes the PTY for the given shell.
// Returns CodeNotFound if no session exists, or CodeBadRequest for invalid dimensions.
func (m *LocalManager) Resize(workspaceID string, shellID string, cols, rows int) error {
	m.sessionsMu.RLock()
	defer m.sessionsMu.RUnlock()

	ws, ok := m.sessions[workspaceID]
	if !ok {
		return util.New(util.CodeNotFound, "shell not found")
	}
	sess, ok := ws[shellID]
	if !ok {
		return util.New(util.CodeNotFound, "shell not found")
	}
	if cols <= 0 || rows <= 0 || cols > 10000 || rows > 10000 {
		return util.New(util.CodeBadRequest, "invalid cols/rows")
	}
	f, ok := sess.Stdin.(*os.File)
	if !ok {
		return util.New(util.CodeInternal, "shell stdin is not a pty file")
	}
	return pty.Setsize(f, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows), X: 0, Y: 0})
}

// Status returns the current shell status for the workspace.
func (m *LocalManager) Status(workspaceID string, shellID string) Status {
	m.sessionsMu.RLock()
	ws, ok := m.sessions[workspaceID]
	if !ok {
		m.sessionsMu.RUnlock()
		return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false}
	}
	sess, ok := ws[shellID]
	m.sessionsMu.RUnlock()
	if !ok {
		return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false}
	}

	if m.IsExited(shellID) {
		return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false}
	}

	return Status{
		WorkspaceID: workspaceID,
		ShellID:     shellID,
		Running:     true,
		PID:         sess.Cmd.Process.Pid,
	}
}

// Subscribe creates a new output subscription for the identified shell.
// Returns CodeNotFound if the shell has no session.
func (m *LocalManager) Subscribe(workspaceID string, shellID string) (Subscription, error) {
	m.sessionsMu.RLock()
	ws, ok := m.sessions[workspaceID]
	if !ok {
		m.sessionsMu.RUnlock()
		return nil, util.New(util.CodeNotFound, "shell not found")
	}
	sess, ok := ws[shellID]
	m.sessionsMu.RUnlock()
	if !ok {
		return nil, util.New(util.CodeNotFound, "shell not found")
	}

	sess.mu.Lock()
	sub := newSubscription(sess)
	sess.Subscribers[sub] = struct{}{}
	sess.mu.Unlock()

	if m.IsExited(shellID) {
		sub.closeChan()
	}

	return sub, nil
}

// Stdin returns the stdin writer for the identified shell.
func (m *LocalManager) Stdin(workspaceID string, shellID string) (io.WriteCloser, error) {
	return m.stdinOf(workspaceID, shellID)
}

// ExitStatus returns the exit status for the identified shell.
func (m *LocalManager) ExitStatus(workspaceID string, shellID string) (ExitInfo, error) {
	return m.exitOf(workspaceID, shellID)
}

// List returns all shells for the workspace in creation order.
func (m *LocalManager) List(workspaceID string) []ShellInfo {
	m.sessionsMu.Lock()
	defer m.sessionsMu.Unlock()

	var infos []ShellInfo
	for _, shellID := range m.shellOrder[workspaceID] {
		ws, ok := m.sessions[workspaceID]
		if !ok {
			continue
		}
		sess, ok := ws[shellID]
		if !ok {
			continue
		}

		info := ShellInfo{
			ShellID:     shellID,
			WorkspaceID: workspaceID,
			CreatedAt:   sess.CreatedAt,
		}

		if !m.IsExited(shellID) && sess.Cmd.Process != nil {
			info.Running = true
			info.PID = sess.Cmd.Process.Pid
		}

		sess.mu.Lock()
		info.ExitCode = sess.Exit.Code
		info.Signal = sess.Exit.Signal
		sess.mu.Unlock()

		infos = append(infos, info)
	}
	return infos
}

// stdinOf returns the session's stdin writer or CodeNotFound.
func (m *LocalManager) stdinOf(workspaceID string, shellID string) (io.WriteCloser, error) {
	m.sessionsMu.RLock()
	defer m.sessionsMu.RUnlock()

	ws, ok := m.sessions[workspaceID]
	if !ok {
		return nil, util.New(util.CodeNotFound, "shell not found")
	}
	sess, ok := ws[shellID]
	if !ok {
		return nil, util.New(util.CodeNotFound, "shell not found")
	}
	return sess.Stdin, nil
}

// exitOf returns the session's recorded exit info or CodeNotFound.
func (m *LocalManager) exitOf(workspaceID string, shellID string) (ExitInfo, error) {
	m.sessionsMu.RLock()
	ws, ok := m.sessions[workspaceID]
	if !ok {
		m.sessionsMu.RUnlock()
		return ExitInfo{}, util.New(util.CodeNotFound, "shell not found")
	}
	sess, ok := ws[shellID]
	m.sessionsMu.RUnlock()
	if !ok {
		return ExitInfo{}, util.New(util.CodeNotFound, "shell not found")
	}

	sess.mu.Lock()
	defer sess.mu.Unlock()
	return sess.Exit, nil
}