codespace/internal/shell/manager.go

package shell

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

	"codespace/internal/util"

	"github.com/creack/pty"
)

const (
	outputBufferSize = 4 * 1024
)

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

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

// LocalManager implements Manager using local OS processes.
type LocalManager struct {
	command  string
	args     []string
	mu       sync.Mutex
	sessions map[string]*Session
	exitedMu sync.Mutex
	exited   map[string]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]*Session),
		exited:   make(map[string]struct{}),
	}
}

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

// Start launches a shell for the given workspace.
// Returns CodeConflict if a session is already running.
func (m *LocalManager) Start(workspaceID string, workspaceRoot string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.sessions[workspaceID]; ok {
		m.exitedMu.Lock()
		_, exited := m.exited[workspaceID]
		m.exitedMu.Unlock()
		if !exited {
			return util.New(util.CodeConflict, "shell already running")
		}
	}

	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{
		WorkspaceID: workspaceID,
		Root:        workspaceRoot,
		Cmd:         cmd,
		Stdin:       ptyF,
		Subscribers: make(map[Subscription]struct{}),
	}
	m.sessions[workspaceID] = sess

	m.exitedMu.Lock()
	delete(m.exited, workspaceID)
	m.exitedMu.Unlock()

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

	return 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.exitedMu.Lock()
	m.exited[sess.WorkspaceID] = struct{}{}
	m.exitedMu.Unlock()

	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.
// Returns CodeNotFound if no session exists.
func (m *LocalManager) Stop(workspaceID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	sess, ok := m.sessions[workspaceID]
	if !ok {
		return util.New(util.CodeNotFound, "no running shell for workspace")
	}

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

	delete(m.sessions, workspaceID)
	m.exitedMu.Lock()
	delete(m.exited, workspaceID)
	m.exitedMu.Unlock()
	return nil
}

// Restart stops (if running) then starts the shell.
func (m *LocalManager) Restart(workspaceID string, workspaceRoot string) error {
	m.mu.Lock()
	if sess, ok := m.sessions[workspaceID]; ok {
		m.exitedMu.Lock()
		_, exited := m.exited[workspaceID]
		m.exitedMu.Unlock()
		if !exited && sess.Cmd.Process != nil {
			_ = sess.Cmd.Process.Kill()
		}
		delete(m.sessions, workspaceID)
		m.exitedMu.Lock()
		delete(m.exited, workspaceID)
		m.exitedMu.Unlock()
	}
	m.mu.Unlock()

	return m.Start(workspaceID, workspaceRoot)
}

// Resize resizes the PTY for the given workspace.
// Returns CodeNotFound if no session exists, or CodeBadRequest for invalid dimensions.
func (m *LocalManager) Resize(workspaceID string, cols, rows int) error {
	m.mu.Lock()
	sess, ok := m.sessions[workspaceID]
	m.mu.Unlock()
	if !ok {
		return util.New(util.CodeNotFound, "no running shell for workspace")
	}
	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) Status {
	m.mu.Lock()
	sess, ok := m.sessions[workspaceID]
	m.mu.Unlock()

	if !ok {
		return Status{WorkspaceID: workspaceID, Running: false}
	}

	m.exitedMu.Lock()
	_, exited := m.exited[workspaceID]
	m.exitedMu.Unlock()
	if exited {
		return Status{WorkspaceID: workspaceID, Running: false}
	}

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

// Subscribe creates a new output subscription for the workspace.
// Returns CodeNotFound if the workspace has no session.
func (m *LocalManager) Subscribe(workspaceID string) (Subscription, error) {
	m.mu.Lock()
	sess, ok := m.sessions[workspaceID]
	m.mu.Unlock()

	if !ok {
		return nil, util.New(util.CodeNotFound, "no running shell")
	}

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

	m.exitedMu.Lock()
	_, exited := m.exited[workspaceID]
	m.exitedMu.Unlock()
	if exited {
		sub.closeChan()
	}

	return sub, nil
}

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

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

// stdinOf returns the session's stdin writer or CodeNotFound.
func (m *LocalManager) stdinOf(workspaceID string) (io.WriteCloser, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	sess, ok := m.sessions[workspaceID]
	if !ok {
		return nil, util.New(util.CodeNotFound, "no running shell")
	}
	return sess.Stdin, nil
}

// exitOf returns the session's recorded exit info or CodeNotFound.
func (m *LocalManager) exitOf(workspaceID string) (ExitInfo, error) {
	m.mu.Lock()
	sess, ok := m.sessions[workspaceID]
	m.mu.Unlock()
	if !ok {
		return ExitInfo{}, util.New(util.CodeNotFound, "no running shell")
	}

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