codespace/internal/process/manager.go

package process

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

	"codespace/internal/util"
)

const (
	outputBufferSize = 4 * 1024
)

// Manager manages OpenCode processes 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.
//
// Concurrency model:
//   - sessionsMu is an RWMutex. Read-only paths (Status, Subscribe lookup,
//     Stdin lookup, ExitStatus lookup) take RLock; mutating paths (Start,
//     Stop, Restart) take Lock.
//   - exited is a sync.Map. waitExit is the only goroutine that calls
//     sess.Cmd.Wait() and therefore the only one that writes here. Reads
//     from anywhere (Start, Status, Subscribe, Restart) are lock-free.
type LocalManager struct {
	command    string
	args       []string
	sessionsMu sync.RWMutex
	sessions   map[string]*Session
	exited     sync.Map
}

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

// IsExited reports whether the process for the workspace has exited.
// Lock-free; safe to call from any goroutine.
func (m *LocalManager) IsExited(workspaceID string) bool {
	_, exited := m.exited.Load(workspaceID)
	return exited
}

// MarkAsExited records that the process for the workspace has exited.
// Called only from waitExit after sess.Cmd.Wait() returns.
func (m *LocalManager) MarkAsExited(workspaceID string) {
	m.exited.Store(workspaceID, struct{}{})
}

// ClearExited removes the exited marker for a workspace. Called when
// (re)starting a process so the next Status/Subscribe call won't see
// a stale exit.
func (m *LocalManager) ClearExited(workspaceID string) {
	m.exited.Delete(workspaceID)
}

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

	if _, ok := m.sessions[workspaceID]; ok && !m.IsExited(workspaceID) {
		return util.New(util.CodeConflict, "process 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 process", err)
	}

	stdinR.Close()
	stdoutW.Close()

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

	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 process to exit, then records the exit status, closes
// the stdin writer, and closes every subscriber channel exactly once.
//
// This is the only goroutine that calls sess.Cmd.Wait() and the only one
// that writes to m.exited. All other code paths use m.IsExited instead of
// reading sess.Cmd.ProcessState directly (which would race with Wait).
func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) {
	_ = sess.Cmd.Wait()
	<-outputDone

	m.MarkAsExited(sess.WorkspaceID)

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

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

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

	delete(m.sessions, workspaceID)
	m.ClearExited(workspaceID)
	return nil
}

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

	return m.Start(workspaceID, workspaceRoot)
}

// Status returns the current process status for the workspace.
func (m *LocalManager) Status(workspaceID string) Status {
	m.sessionsMu.RLock()
	sess, ok := m.sessions[workspaceID]
	m.sessionsMu.RUnlock()

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

	if m.IsExited(workspaceID) {
		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.sessionsMu.RLock()
	sess, ok := m.sessions[workspaceID]
	m.sessionsMu.RUnlock()

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

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

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

	return sub, nil
}

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

// ExitStatus returns the exit status for the workspace process.
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.sessionsMu.RLock()
	defer m.sessionsMu.RUnlock()
	sess, ok := m.sessions[workspaceID]
	if !ok {
		return nil, util.New(util.CodeNotFound, "no running process")
	}
	return sess.Stdin, nil
}

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

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