Carlos Florêncio
673 lines
16 KiB
Go
673 lines
16 KiB
Go
// Package shell provides an embedded PTY-based shell session for agentctl.
|
|
// The shell session is automatically created when agentctl starts and destroyed when it stops.
|
|
package shell
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"os/exec"
|
|
"runtime"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/creack/pty"
|
|
"github.com/kandev/kandev/internal/common/logger"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
// Session represents the single embedded shell session for an agentctl instance.
|
|
// It is automatically created when agentctl starts and destroyed when it stops.
|
|
type Session struct {
|
|
logger *logger.Logger
|
|
workDir string
|
|
shell string // Shell command (detected based on OS)
|
|
shellArgs []string // Shell arguments
|
|
config Config // Stored config for respawn
|
|
|
|
// PTY and process
|
|
pty *os.File
|
|
cmd *exec.Cmd
|
|
lifecycle shellProcessLifecycleHandle
|
|
reapErr error
|
|
|
|
// State
|
|
running bool
|
|
startedAt time.Time
|
|
mu sync.RWMutex
|
|
|
|
// Output subscribers (WebSocket connections)
|
|
subscribers map[chan<- []byte]struct{}
|
|
subMu sync.RWMutex
|
|
|
|
// Ring buffer for recent output (for new subscribers)
|
|
outputBuffer []byte
|
|
bufferMu sync.RWMutex
|
|
|
|
// Lifecycle
|
|
stopCh chan struct{}
|
|
doneCh chan struct{}
|
|
stopping bool // true when Stop() has been called (don't respawn)
|
|
|
|
killGroupFn func(*os.Process) error
|
|
waitGroupExitFn func(*os.Process) bool
|
|
beforeRespawn func()
|
|
afterStopClaim func()
|
|
}
|
|
|
|
// Maximum size of output buffer (16KB should be enough for recent history)
|
|
const maxOutputBufferSize = 16 * 1024
|
|
const shellStopGrace = 250 * time.Millisecond
|
|
const shellReapGrace = 2 * time.Second
|
|
|
|
func killAndWaitShellCommand(cmd *exec.Cmd) error {
|
|
if cmd == nil || cmd.Process == nil {
|
|
return nil
|
|
}
|
|
_ = cmd.Process.Kill()
|
|
done := make(chan struct{})
|
|
go func() {
|
|
_ = cmd.Wait()
|
|
close(done)
|
|
}()
|
|
select {
|
|
case <-done:
|
|
return nil
|
|
case <-time.After(shellReapGrace):
|
|
return fmt.Errorf("shell process %d was not reaped after failed startup", cmd.Process.Pid)
|
|
}
|
|
}
|
|
|
|
// Config holds shell session configuration
|
|
type Config struct {
|
|
WorkDir string // Working directory (default: workspace)
|
|
Cols int // Initial terminal columns (default: 80)
|
|
Rows int // Initial terminal rows (default: 24)
|
|
ShellCommand string // Optional shell command override
|
|
ShellArgs []string // Optional shell args override
|
|
Env map[string]string // Optional environment overrides
|
|
}
|
|
|
|
// DefaultConfig returns the default shell configuration
|
|
func DefaultConfig(workDir string) Config {
|
|
return Config{
|
|
WorkDir: workDir,
|
|
Cols: 80,
|
|
Rows: 24,
|
|
}
|
|
}
|
|
|
|
// Status represents the shell status
|
|
type Status struct {
|
|
Running bool
|
|
Pid int
|
|
Shell string
|
|
Cwd string
|
|
StartedAt time.Time
|
|
}
|
|
|
|
// detectShell returns the appropriate shell for the current OS.
|
|
// It validates that the shell exists before returning it.
|
|
func detectShell() (string, []string) {
|
|
if runtime.GOOS == "windows" {
|
|
// Prefer PowerShell if available
|
|
if _, err := exec.LookPath("pwsh.exe"); err == nil {
|
|
return "pwsh.exe", []string{"-NoLogo", "-NoExit"}
|
|
}
|
|
if _, err := exec.LookPath("powershell.exe"); err == nil {
|
|
return "powershell.exe", []string{"-NoLogo", "-NoExit"}
|
|
}
|
|
return "cmd.exe", nil
|
|
}
|
|
|
|
// Unix-like systems (Linux, macOS)
|
|
// Check $SHELL but only use it if the shell actually exists
|
|
// (host's $SHELL may not be installed in Docker containers).
|
|
// No -l flag: Debian's /etc/profile (and /etc/profile.d/*) can overwrite
|
|
// PATH, dropping container-set entries like /data/.npm-global/bin where
|
|
// agent CLIs (claude, codex, ...) are installed.
|
|
if shell := os.Getenv("SHELL"); shell != "" {
|
|
if _, err := exec.LookPath(shell); err == nil {
|
|
return shell, nil
|
|
}
|
|
}
|
|
|
|
// Try common shells
|
|
shells := []string{"/bin/bash", "/bin/zsh", "/bin/sh"}
|
|
for _, sh := range shells {
|
|
if _, err := os.Stat(sh); err == nil {
|
|
return sh, nil
|
|
}
|
|
}
|
|
|
|
return "/bin/sh", nil
|
|
}
|
|
|
|
// NewSession creates and starts a new shell session.
|
|
// This is called automatically when agentctl starts.
|
|
func NewSession(cfg Config, log *logger.Logger) (*Session, error) {
|
|
shell, args := detectShell()
|
|
if cfg.ShellCommand != "" {
|
|
shell = cfg.ShellCommand
|
|
if len(cfg.ShellArgs) > 0 {
|
|
args = cfg.ShellArgs
|
|
} else {
|
|
args = defaultShellArgs(shell)
|
|
}
|
|
}
|
|
|
|
s := &Session{
|
|
logger: log.WithFields(zap.String("component", "shell")),
|
|
workDir: cfg.WorkDir,
|
|
shell: shell,
|
|
shellArgs: args,
|
|
config: cfg,
|
|
subscribers: make(map[chan<- []byte]struct{}),
|
|
stopCh: make(chan struct{}),
|
|
doneCh: make(chan struct{}),
|
|
}
|
|
|
|
if err := s.start(cfg); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return s, nil
|
|
}
|
|
|
|
func defaultShellArgs(shellCommand string) []string {
|
|
if runtime.GOOS == "windows" {
|
|
lower := strings.ToLower(shellCommand)
|
|
if strings.Contains(lower, "pwsh") || strings.Contains(lower, "powershell") {
|
|
return []string{"-NoLogo", "-NoExit"}
|
|
}
|
|
return nil
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// start initializes and starts the shell process
|
|
func (s *Session) start(cfg Config) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
doneCh := make(chan struct{})
|
|
|
|
s.cmd = exec.Command(s.shell, s.shellArgs...)
|
|
s.cmd.Dir = cfg.WorkDir
|
|
s.cmd.Env = buildShellEnv(cfg.WorkDir, cfg.Env)
|
|
configureShellProcess(s.cmd)
|
|
|
|
// Start with PTY
|
|
var err error
|
|
s.pty, err = pty.StartWithSize(s.cmd, &pty.Winsize{
|
|
Cols: uint16(cfg.Cols),
|
|
Rows: uint16(cfg.Rows),
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("failed to start PTY: %w", err)
|
|
}
|
|
lifecycle, lifecycleErr := installShellProcessLifecycle(s.cmd)
|
|
if lifecycleErr != nil {
|
|
_ = s.pty.Close()
|
|
reapErr := killAndWaitShellCommand(s.cmd)
|
|
return errors.Join(fmt.Errorf("failed to install shell process lifecycle: %w", lifecycleErr), reapErr)
|
|
}
|
|
|
|
s.running = true
|
|
s.startedAt = time.Now()
|
|
s.doneCh = doneCh
|
|
s.lifecycle = lifecycle
|
|
s.reapErr = nil
|
|
|
|
s.logger.Info("shell session started",
|
|
zap.String("shell", s.shell),
|
|
zap.String("cwd", cfg.WorkDir),
|
|
zap.Int("pid", s.cmd.Process.Pid))
|
|
|
|
// Start output reader
|
|
go s.readOutput()
|
|
|
|
// Wait for process exit
|
|
go s.waitForExit(s.cmd, doneCh, lifecycle)
|
|
|
|
return nil
|
|
}
|
|
|
|
// Stop gracefully stops the shell session.
|
|
// Called automatically when agentctl stops.
|
|
func (s *Session) Stop() error {
|
|
s.mu.Lock()
|
|
cmd := s.cmd
|
|
doneCh := s.doneCh
|
|
var process *os.Process
|
|
if cmd != nil {
|
|
process = cmd.Process
|
|
}
|
|
pid := 0
|
|
if process != nil {
|
|
pid = process.Pid
|
|
}
|
|
lifecycle := s.lifecycle
|
|
if !s.running {
|
|
s.stopping = true
|
|
doneCh := s.doneCh
|
|
s.mu.Unlock()
|
|
if s.afterStopClaim != nil {
|
|
s.afterStopClaim()
|
|
}
|
|
if doneCh != nil {
|
|
select {
|
|
case <-doneCh:
|
|
return s.cleanupShellProcessGroup(process, pid, "stop_retry", lifecycle)
|
|
case <-time.After(shellReapGrace):
|
|
return fmt.Errorf("shell process was not reaped after stop")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
s.running = false
|
|
s.stopping = true // Mark as stopping to prevent respawn
|
|
s.mu.Unlock()
|
|
|
|
s.logger.Info("stopping shell session",
|
|
zap.String("shell", s.shell),
|
|
zap.String("cwd", s.workDir),
|
|
zap.Int("pid", pid),
|
|
zap.Duration("grace", shellStopGrace))
|
|
s.logger.Debug("shell session stop requested",
|
|
zap.Int("pid", pid),
|
|
zap.String("shell", s.shell),
|
|
zap.String("cwd", s.workDir))
|
|
|
|
// Signal stop
|
|
close(s.stopCh)
|
|
|
|
// Close PTY (sends SIGHUP to process on Unix)
|
|
if s.pty != nil {
|
|
s.logger.Debug("shell session PTY close requested", zap.Int("pid", pid))
|
|
_ = s.pty.Close()
|
|
}
|
|
|
|
// Wait for exit with timeout
|
|
select {
|
|
case <-doneCh:
|
|
s.logger.Info("shell session stopped gracefully")
|
|
return s.cleanupShellProcessGroup(process, pid, "leader_exited", lifecycle)
|
|
case <-time.After(shellStopGrace):
|
|
s.logger.Warn("shell session stop timeout, force killing",
|
|
zap.Int("pid", pid),
|
|
zap.Duration("grace", shellStopGrace))
|
|
if err := s.cleanupShellProcessGroup(process, pid, "stop_timeout", lifecycle); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if pid == 0 {
|
|
return nil
|
|
}
|
|
|
|
select {
|
|
case <-doneCh:
|
|
return nil
|
|
case <-time.After(shellReapGrace):
|
|
return fmt.Errorf("shell process %d was not reaped after force kill", pid)
|
|
}
|
|
}
|
|
|
|
func (s *Session) cleanupShellProcessGroup(
|
|
process *os.Process,
|
|
pid int,
|
|
reason string,
|
|
lifecycle shellProcessLifecycleHandle,
|
|
) error {
|
|
if process == nil {
|
|
return nil
|
|
}
|
|
s.logger.Debug("shell session process group cleanup requested",
|
|
zap.Int("pid", pid),
|
|
zap.String("reason", reason))
|
|
if ownsShellProcessLifecycle(lifecycle) {
|
|
if err := reapShellProcessLifecycle(lifecycle); err != nil {
|
|
return fmt.Errorf("reap shell process job: %w", err)
|
|
}
|
|
s.mu.Lock()
|
|
if s.lifecycle == lifecycle {
|
|
s.lifecycle = shellProcessLifecycleHandle{}
|
|
s.reapErr = nil
|
|
}
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
killGroup := killShellProcessGroup
|
|
if s.killGroupFn != nil {
|
|
killGroup = s.killGroupFn
|
|
}
|
|
if err := killGroup(process); err != nil {
|
|
s.logger.Debug("shell session process group cleanup failed",
|
|
zap.Int("pid", pid),
|
|
zap.String("reason", reason),
|
|
zap.Error(err))
|
|
return err
|
|
}
|
|
if s.waitGroupExitFn != nil {
|
|
if !s.waitGroupExitFn(process) {
|
|
return fmt.Errorf("shell process group %d remains alive after force kill", pid)
|
|
}
|
|
return nil
|
|
}
|
|
deadline := time.NewTimer(shellReapGrace)
|
|
defer deadline.Stop()
|
|
ticker := time.NewTicker(20 * time.Millisecond)
|
|
defer ticker.Stop()
|
|
for shellProcessGroupAlive(process) {
|
|
select {
|
|
case <-deadline.C:
|
|
return fmt.Errorf("shell process group %d remains alive after force kill", pid)
|
|
case <-ticker.C:
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Write sends input to the shell
|
|
func (s *Session) Write(data []byte) (int, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
if !s.running || s.pty == nil {
|
|
return 0, fmt.Errorf("shell not running")
|
|
}
|
|
|
|
return s.pty.Write(data)
|
|
}
|
|
|
|
// Resize changes the terminal dimensions of the shell PTY.
|
|
func (s *Session) Resize(cols, rows uint16) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if !s.running || s.pty == nil {
|
|
return fmt.Errorf("shell not running")
|
|
}
|
|
|
|
// Store for respawn
|
|
s.config.Cols = int(cols)
|
|
s.config.Rows = int(rows)
|
|
|
|
return pty.Setsize(s.pty, &pty.Winsize{Cols: cols, Rows: rows})
|
|
}
|
|
|
|
// Subscribe adds a subscriber for shell output
|
|
func (s *Session) Subscribe(ch chan<- []byte) {
|
|
s.subMu.Lock()
|
|
defer s.subMu.Unlock()
|
|
s.subscribers[ch] = struct{}{}
|
|
}
|
|
|
|
// Unsubscribe removes a subscriber
|
|
func (s *Session) Unsubscribe(ch chan<- []byte) {
|
|
s.subMu.Lock()
|
|
defer s.subMu.Unlock()
|
|
delete(s.subscribers, ch)
|
|
}
|
|
|
|
// Status returns the current shell status
|
|
func (s *Session) Status() Status {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
pid := 0
|
|
if s.cmd != nil && s.cmd.Process != nil {
|
|
pid = s.cmd.Process.Pid
|
|
}
|
|
|
|
return Status{
|
|
Running: s.running,
|
|
Pid: pid,
|
|
Shell: s.shell,
|
|
Cwd: s.workDir,
|
|
StartedAt: s.startedAt,
|
|
}
|
|
}
|
|
|
|
// Config returns a defensive copy of the session configuration.
|
|
func (s *Session) Config() Config {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
cfg := s.config
|
|
cfg.ShellArgs = append([]string(nil), s.config.ShellArgs...)
|
|
if s.config.Env != nil {
|
|
cfg.Env = make(map[string]string, len(s.config.Env))
|
|
for key, value := range s.config.Env {
|
|
cfg.Env[key] = value
|
|
}
|
|
}
|
|
return cfg
|
|
}
|
|
|
|
// readOutput continuously reads from PTY and broadcasts to subscribers
|
|
func (s *Session) readOutput() {
|
|
buf := make([]byte, 4096)
|
|
|
|
for {
|
|
select {
|
|
case <-s.stopCh:
|
|
return
|
|
default:
|
|
}
|
|
|
|
n, err := s.pty.Read(buf)
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
s.logger.Debug("shell read error", zap.Error(err))
|
|
}
|
|
return
|
|
}
|
|
|
|
if n > 0 {
|
|
data := make([]byte, n)
|
|
copy(data, buf[:n])
|
|
s.broadcast(data)
|
|
}
|
|
}
|
|
}
|
|
|
|
// broadcast sends data to all subscribers and stores in buffer
|
|
func (s *Session) broadcast(data []byte) {
|
|
// Store in ring buffer for new subscribers
|
|
s.appendToBuffer(data)
|
|
|
|
s.subMu.RLock()
|
|
defer s.subMu.RUnlock()
|
|
|
|
for ch := range s.subscribers {
|
|
select {
|
|
case ch <- data:
|
|
default:
|
|
// Subscriber channel full, skip
|
|
}
|
|
}
|
|
}
|
|
|
|
// appendToBuffer adds data to the output ring buffer
|
|
func (s *Session) appendToBuffer(data []byte) {
|
|
s.bufferMu.Lock()
|
|
defer s.bufferMu.Unlock()
|
|
|
|
s.outputBuffer = append(s.outputBuffer, data...)
|
|
|
|
// Trim if exceeds max size (keep the most recent data)
|
|
if len(s.outputBuffer) > maxOutputBufferSize {
|
|
s.outputBuffer = s.outputBuffer[len(s.outputBuffer)-maxOutputBufferSize:]
|
|
}
|
|
}
|
|
|
|
// GetBufferedOutput returns the buffered recent output
|
|
func (s *Session) GetBufferedOutput() []byte {
|
|
s.bufferMu.RLock()
|
|
defer s.bufferMu.RUnlock()
|
|
|
|
if len(s.outputBuffer) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Return a copy to avoid race conditions
|
|
result := make([]byte, len(s.outputBuffer))
|
|
copy(result, s.outputBuffer)
|
|
return result
|
|
}
|
|
|
|
// waitForExit waits for the shell process to exit and respawns if not stopping
|
|
func (s *Session) waitForExit(cmd *exec.Cmd, doneCh chan struct{}, lifecycle shellProcessLifecycleHandle) {
|
|
if cmd != nil {
|
|
_ = cmd.Wait()
|
|
}
|
|
pid := 0
|
|
var process *os.Process
|
|
if cmd != nil {
|
|
process = cmd.Process
|
|
}
|
|
if process != nil {
|
|
pid = process.Pid
|
|
}
|
|
reapErr := s.cleanupShellProcessGroup(process, pid, "leader_exited", lifecycle)
|
|
|
|
s.mu.Lock()
|
|
stopping := s.stopping
|
|
s.running = false
|
|
s.reapErr = reapErr
|
|
s.mu.Unlock()
|
|
if reapErr != nil {
|
|
s.logger.Error("shell process tree was not reaped", zap.Int("pid", pid), zap.Error(reapErr))
|
|
close(doneCh)
|
|
return
|
|
}
|
|
|
|
// If Stop() was called, don't respawn - just signal done
|
|
if stopping {
|
|
s.logger.Info("shell process exited (stopping)")
|
|
close(doneCh)
|
|
return
|
|
}
|
|
|
|
s.logger.Info("shell process exited unexpectedly, respawning...")
|
|
|
|
// Small delay before respawn to avoid tight loop on repeated failures.
|
|
if s.beforeRespawn != nil {
|
|
s.beforeRespawn()
|
|
} else {
|
|
time.Sleep(100 * time.Millisecond)
|
|
}
|
|
|
|
// Respawn the shell
|
|
if err := s.respawn(); err != nil {
|
|
s.logger.Error("failed to respawn shell", zap.Error(err))
|
|
close(doneCh)
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
stoppedBeforeRespawn := s.stopping && !s.running
|
|
s.mu.Unlock()
|
|
if stoppedBeforeRespawn {
|
|
close(doneCh)
|
|
}
|
|
}
|
|
|
|
// respawn restarts the shell process after an unexpected exit
|
|
func (s *Session) respawn() error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
// Check if we're stopping (Stop() called during respawn)
|
|
if s.stopping {
|
|
return nil
|
|
}
|
|
|
|
// Close old PTY to prevent file descriptor leak.
|
|
// Don't nil out s.pty — readOutput() may still be reading from it
|
|
// concurrently without the mutex. Closing is sufficient: any in-flight
|
|
// Read() will return an error and the goroutine will exit.
|
|
if s.pty != nil {
|
|
_ = s.pty.Close()
|
|
}
|
|
|
|
doneCh := make(chan struct{})
|
|
s.cmd = exec.Command(s.shell, s.shellArgs...)
|
|
s.cmd.Dir = s.workDir
|
|
s.cmd.Env = buildShellEnv(s.workDir, s.config.Env)
|
|
configureShellProcess(s.cmd)
|
|
|
|
// Start with PTY
|
|
var err error
|
|
s.pty, err = pty.StartWithSize(s.cmd, &pty.Winsize{
|
|
Cols: uint16(s.config.Cols),
|
|
Rows: uint16(s.config.Rows),
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("failed to start PTY: %w", err)
|
|
}
|
|
lifecycle, lifecycleErr := installShellProcessLifecycle(s.cmd)
|
|
if lifecycleErr != nil {
|
|
_ = s.pty.Close()
|
|
reapErr := killAndWaitShellCommand(s.cmd)
|
|
return errors.Join(fmt.Errorf("failed to install shell process lifecycle: %w", lifecycleErr), reapErr)
|
|
}
|
|
|
|
s.running = true
|
|
s.startedAt = time.Now()
|
|
s.doneCh = doneCh
|
|
s.lifecycle = lifecycle
|
|
s.reapErr = nil
|
|
|
|
s.logger.Info("shell session respawned",
|
|
zap.String("shell", s.shell),
|
|
zap.String("cwd", s.workDir),
|
|
zap.Int("pid", s.cmd.Process.Pid))
|
|
|
|
// Start output reader
|
|
go s.readOutput()
|
|
|
|
// Wait for process exit (recursive respawn on exit)
|
|
go s.waitForExit(s.cmd, doneCh, lifecycle)
|
|
|
|
return nil
|
|
}
|
|
|
|
// buildShellEnv creates the environment for the shell process
|
|
func buildShellEnv(workDir string, overrides map[string]string) []string {
|
|
env := os.Environ()
|
|
for key, value := range overrides {
|
|
env = upsertShellEnv(env, key, value)
|
|
}
|
|
|
|
// Set working directory related vars
|
|
env = append(env, "PWD="+workDir)
|
|
|
|
// Set terminal type
|
|
env = append(env, "TERM=xterm-256color")
|
|
env = append(env, "LANG=C.UTF-8")
|
|
env = append(env, "LC_ALL=C.UTF-8")
|
|
|
|
return env
|
|
}
|
|
|
|
func upsertShellEnv(env []string, key, value string) []string {
|
|
prefix := key + "="
|
|
result := make([]string, 0, len(env)+1)
|
|
found := false
|
|
for _, entry := range env {
|
|
if strings.HasPrefix(entry, prefix) {
|
|
if !found {
|
|
result = append(result, prefix+value)
|
|
found = true
|
|
}
|
|
continue
|
|
}
|
|
result = append(result, entry)
|
|
}
|
|
if !found {
|
|
result = append(result, prefix+value)
|
|
}
|
|
return result
|
|
}
|