background.go

  1package shell
  2
  3import (
  4	"bytes"
  5	"context"
  6	"fmt"
  7	"sync"
  8
  9	"github.com/charmbracelet/hotdiva2000"
 10)
 11
 12// BackgroundShell represents a shell running in the background.
 13type BackgroundShell struct {
 14	ID         string
 15	Shell      *Shell
 16	ctx        context.Context
 17	cancel     context.CancelFunc
 18	stdout     *bytes.Buffer
 19	stderr     *bytes.Buffer
 20	mu         sync.RWMutex
 21	done       chan struct{}
 22	exitErr    error
 23	workingDir string
 24}
 25
 26// BackgroundShellManager manages background shell instances.
 27type BackgroundShellManager struct {
 28	shells map[string]*BackgroundShell
 29	mu     sync.RWMutex
 30}
 31
 32var (
 33	backgroundManager     *BackgroundShellManager
 34	backgroundManagerOnce sync.Once
 35)
 36
 37// GetBackgroundShellManager returns the singleton background shell manager.
 38func GetBackgroundShellManager() *BackgroundShellManager {
 39	backgroundManagerOnce.Do(func() {
 40		backgroundManager = &BackgroundShellManager{
 41			shells: make(map[string]*BackgroundShell),
 42		}
 43	})
 44	return backgroundManager
 45}
 46
 47// Start creates and starts a new background shell with the given command.
 48func (m *BackgroundShellManager) Start(ctx context.Context, workingDir string, blockFuncs []BlockFunc, command string) (*BackgroundShell, error) {
 49	id := hotdiva2000.Generate()
 50
 51	shell := NewShell(&Options{
 52		WorkingDir: workingDir,
 53		BlockFuncs: blockFuncs,
 54	})
 55
 56	shellCtx, cancel := context.WithCancel(ctx)
 57
 58	bgShell := &BackgroundShell{
 59		ID:         id,
 60		Shell:      shell,
 61		ctx:        shellCtx,
 62		cancel:     cancel,
 63		stdout:     &bytes.Buffer{},
 64		stderr:     &bytes.Buffer{},
 65		done:       make(chan struct{}),
 66		workingDir: workingDir,
 67	}
 68
 69	m.mu.Lock()
 70	m.shells[id] = bgShell
 71	m.mu.Unlock()
 72
 73	go func() {
 74		defer close(bgShell.done)
 75
 76		stdout, stderr, err := shell.Exec(shellCtx, command)
 77
 78		bgShell.mu.Lock()
 79		bgShell.stdout.WriteString(stdout)
 80		bgShell.stderr.WriteString(stderr)
 81		bgShell.exitErr = err
 82		bgShell.mu.Unlock()
 83	}()
 84
 85	return bgShell, nil
 86}
 87
 88// Get retrieves a background shell by ID.
 89func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
 90	m.mu.RLock()
 91	defer m.mu.RUnlock()
 92	shell, ok := m.shells[id]
 93	return shell, ok
 94}
 95
 96// Kill terminates a background shell by ID.
 97func (m *BackgroundShellManager) Kill(id string) error {
 98	m.mu.Lock()
 99	shell, ok := m.shells[id]
100	if !ok {
101		m.mu.Unlock()
102		return fmt.Errorf("background shell not found: %s", id)
103	}
104	delete(m.shells, id)
105	m.mu.Unlock()
106
107	shell.cancel()
108	<-shell.done
109	return nil
110}
111
112// List returns all background shell IDs.
113func (m *BackgroundShellManager) List() []string {
114	m.mu.RLock()
115	defer m.mu.RUnlock()
116
117	ids := make([]string, 0, len(m.shells))
118	for id := range m.shells {
119		ids = append(ids, id)
120	}
121	return ids
122}
123
124// KillAll terminates all background shells.
125func (m *BackgroundShellManager) KillAll() {
126	m.mu.Lock()
127	shells := make([]*BackgroundShell, 0, len(m.shells))
128	for _, shell := range m.shells {
129		shells = append(shells, shell)
130	}
131	m.shells = make(map[string]*BackgroundShell)
132	m.mu.Unlock()
133
134	for _, shell := range shells {
135		shell.cancel()
136		<-shell.done
137	}
138}
139
140// GetOutput returns the current output of a background shell.
141func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
142	bs.mu.RLock()
143	defer bs.mu.RUnlock()
144
145	select {
146	case <-bs.done:
147		return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr
148	default:
149		return bs.stdout.String(), bs.stderr.String(), false, nil
150	}
151}
152
153// IsDone checks if the background shell has finished execution.
154func (bs *BackgroundShell) IsDone() bool {
155	select {
156	case <-bs.done:
157		return true
158	default:
159		return false
160	}
161}
162
163// Wait blocks until the background shell completes.
164func (bs *BackgroundShell) Wait() {
165	<-bs.done
166}
167
168// GetWorkingDir returns the current working directory of the background shell.
169func (bs *BackgroundShell) GetWorkingDir() string {
170	bs.mu.RLock()
171	defer bs.mu.RUnlock()
172	return bs.workingDir
173}