codespace/internal/shell/manager.go

package shell

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

	"codespace/internal/util"
)

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
}

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

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 sess, ok := m.sessions[workspaceID]; ok && sess.Cmd.Process != nil {
		if sess.Cmd.ProcessState == nil {
			return util.New(util.CodeConflict, "shell already running")
		}
	}

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

	stdinR, stdinW, err := os.Pipe()
	if err != nil {
		return util.Wrap(util.CodeInternal, "failed to create stdin pipe", err)
	}

	stdoutR, stdoutW, err := os.Pipe()
	if err != nil {
		stdinR.Close()
		stdinW.Close()
		return util.Wrap(util.CodeInternal, "failed to create stdout pipe", err)
	}

	cmd.Stdin = stdinR
	cmd.Stdout = stdoutW
	cmd.Stderr = stdoutW

	if err := cmd.Start(); err != nil {
		stdinR.Close()
		stdinW.Close()
		stdoutR.Close()
		stdoutW.Close()
		return util.Wrap(util.CodeInternal, "failed to start shell", err)
	}

	stdinR.Close()
	stdoutW.Close()

	sess := &Session{
		WorkspaceID: workspaceID,
		Root:        workspaceRoot,
		Cmd:         cmd,
		Stdin:       stdinW,
		Subscribers: make(map[Subscription]struct{}),
	}
	m.sessions[workspaceID] = sess

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

	return nil
}

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

	buf := make([]byte, outputBufferSize)

	for {
		n, err := stdoutR.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 pipe is closing.
			}
			break
		}
	}
}

// waitExit waits for the shell to exit, then records the exit status, closes
// the stdin writer, and closes every subscriber channel exactly once.
func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) {
	_ = sess.Cmd.Wait()
	<-outputDone

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

	if sess.Stdin != nil {
		_ = sess.Stdin.Close()
	}

	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 || sess.Cmd.Process == nil {
		return util.New(util.CodeNotFound, "no running shell for workspace")
	}

	if err := sess.Cmd.Process.Kill(); err != nil {
		return util.Wrap(util.CodeInternal, "failed to stop shell", err)
	}

	delete(m.sessions, workspaceID)
	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 && sess.Cmd.Process != nil {
		if sess.Cmd.ProcessState == nil {
			_ = sess.Cmd.Process.Kill()
			delete(m.sessions, workspaceID)
		}
	}
	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()
	defer m.mu.Unlock()

	sess, ok := m.sessions[workspaceID]
	if !ok || sess.Cmd.Process == nil {
		return Status{WorkspaceID: workspaceID, Running: false}
	}

	if sess.Cmd.ProcessState != nil {
		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{}{}
	if sess.Cmd.ProcessState != nil {
		sub.closeChan()
	}
	sess.mu.Unlock()

	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
}