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.
type LocalManager struct {
	command    string
	args       []string
	mu         sync.Mutex
	sessions   map[string]map[string]*Session // workspaceID -> shellID -> Session
	exitedMu   sync.Mutex
	exited     map[string]map[string]struct{} // workspaceID -> shellID -> exited
	shellOrder map[string][]string            // workspaceID -> shellIDs in creation order
}

// 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),
		exited:     make(map[string]map[string]struct{}),
		shellOrder: make(map[string][]string),
	}
}

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

// Start launches a new shell for the given workspace and returns its shellID.
func (m *LocalManager) Start(workspaceID string, workspaceRoot string) (string, error) {
	m.mu.Lock()
	defer m.mu.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

	if m.exited[workspaceID] == nil {
		m.exited[workspaceID] = make(map[string]struct{})
	}
	delete(m.exited[workspaceID], shellID)

	m.shellOrder[workspaceID] = append(m.shellOrder[workspaceID], 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.exitedMu.Lock()
	if m.exited[sess.WorkspaceID] == nil {
		m.exited[sess.WorkspaceID] = make(map[string]struct{})
	}
	m.exited[sess.WorkspaceID][sess.ShellID] = 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 and shellID.
// Returns CodeNotFound if no session exists.
func (m *LocalManager) Stop(workspaceID string, shellID string) error {
	m.mu.Lock()
	defer m.mu.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.exitedMu.Lock()
	if m.exited[workspaceID] != nil {
		delete(m.exited[workspaceID], shellID)
		if len(m.exited[workspaceID]) == 0 {
			delete(m.exited, workspaceID)
		}
	}
	m.exitedMu.Unlock()
	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.mu.Lock()
	ws, ok := m.sessions[workspaceID]
	if ok {
		if sess, ok := ws[shellID]; ok {
			m.exitedMu.Lock()
			_, exited := m.exited[workspaceID][shellID]
			m.exitedMu.Unlock()
			if !exited && sess.Cmd.Process != nil {
				_ = sess.Cmd.Process.Kill()
			}
			delete(ws, shellID)
			if len(ws) == 0 {
				delete(m.sessions, workspaceID)
			}
			m.exitedMu.Lock()
			if m.exited[workspaceID] != nil {
				delete(m.exited[workspaceID], shellID)
				if len(m.exited[workspaceID]) == 0 {
					delete(m.exited, workspaceID)
				}
			}
			m.exitedMu.Unlock()
		}
	}
	m.mu.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.mu.Lock()
	defer m.mu.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 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.mu.Lock()
	defer m.mu.Unlock()

	ws, ok := m.sessions[workspaceID]
	if !ok {
		return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false}
	}
	sess, ok := ws[shellID]
	if !ok {
		return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false}
	}

	m.exitedMu.Lock()
	_, exited := m.exited[workspaceID][shellID]
	m.exitedMu.Unlock()
	if exited {
		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.mu.Lock()
	defer m.mu.Unlock()

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

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

	m.exitedMu.Lock()
	_, exited := m.exited[workspaceID][shellID]
	m.exitedMu.Unlock()
	if exited {
		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.mu.Lock()
	defer m.mu.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,
		}

		m.exitedMu.Lock()
		_, exited := m.exited[workspaceID][shellID]
		m.exitedMu.Unlock()

		if !exited && 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.mu.Lock()
	defer m.mu.Unlock()

	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.mu.Lock()
	defer m.mu.Unlock()

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

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