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)
}

// 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)
}

// 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
}