git.iko.soy/crush
on main help sign in
revert changes on shell
caeb245Raphael Amorim
internal/llm/tools/shell/comparison_test.go+4−2
compare with main
unifiedsplit
@@ -13,7 +13,8 @@ func TestShellPerformanceComparison(t *testing.T) {
13
13
14 // Test quick command
14 // Test quick command
15 start := time.Now()
15 start := time.Now()
16− stdout, stderr, exitCode, _, err := shell.Exec(t.Context(), "echo 'hello'", 0)
16+ stdout, stderr, err := shell.Exec(t.Context(), "echo 'hello'")
17+ exitCode := ExitCode(err)
17 duration := time.Since(start)
18 duration := time.Since(start)
18
19
19 require.NoError(t, err)
20 require.NoError(t, err)
@@ -32,7 +33,8 @@ func BenchmarkShellPolling(b *testing.B) {
32
33
33 for b.Loop() {
34 for b.Loop() {
34 // Use a short sleep to measure polling overhead
35 // Use a short sleep to measure polling overhead
35− _, _, exitCode, _, err := shell.Exec(b.Context(), "sleep 0.02", 0)
36+ _, _, err := shell.Exec(b.Context(), "sleep 0.02")
37+ exitCode := ExitCode(err)
36 if err != nil || exitCode != 0 {
38 if err != nil || exitCode != 0 {
37 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
39 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
38 }
40 }
internal/llm/tools/shell/shell.go+45−281
@@ -1,323 +1,87 @@
1package shell
1package shell
2
2
3import (
3import (
4+ "bytes"
4 "context"
5 "context"
5 "errors"
6 "errors"
6 "fmt"
7 "fmt"
7 "os"
8 "os"
8− "os/exec"
9− "path/filepath"
10 "strings"
9 "strings"
11 "sync"
10 "sync"
12− "syscall"
11+
13− "time"
12+ "github.com/charmbracelet/crush/internal/logging"
13+ "mvdan.cc/sh/v3/expand"
14+ "mvdan.cc/sh/v3/interp"
15+ "mvdan.cc/sh/v3/syntax"
14)
16)
15
17
16type PersistentShell struct {
18type PersistentShell struct {
17− cmd *exec.Cmd
19+ env []string
18− stdin *os.File
19− isAlive bool
20− cwd string
20+ cwd string
21− mu sync.Mutex
21+ mu sync.Mutex
22− commandQueue chan *commandExecution
23−}
24−
25−type commandExecution struct {
26− command string
27− timeout time.Duration
28− resultChan chan commandResult
29− ctx context.Context
30−}
31−
32−type commandResult struct {
33− stdout string
34− stderr string
35− exitCode int
36− interrupted bool
37− err error
38}
22}
39
23
40var (
24var (
25+ once sync.Once
41− shellInstance *PersistentShell
26+ shellInstance *PersistentShell
42− shellInstanceOnce sync.Once
43)
27)
44
28
45−func GetPersistentShell(workingDir string) *PersistentShell {
29+func GetPersistentShell(cwd string) *PersistentShell {
46− shellInstanceOnce.Do(func() {
30+ once.Do(func() {
47− shellInstance = newPersistentShell(workingDir)
31+ shellInstance = newPersistentShell(cwd)
48 })
32 })
49−
50− if shellInstance == nil {
51− shellInstance = newPersistentShell(workingDir)
52− } else if !shellInstance.isAlive {
53− shellInstance = newPersistentShell(shellInstance.cwd)
54− }
55−
56 return shellInstance
33 return shellInstance
57}
34}
58
35
59func newPersistentShell(cwd string) *PersistentShell {
36func newPersistentShell(cwd string) *PersistentShell {
60− // Default to environment variable
61− shellPath := os.Getenv("SHELL")
62− if shellPath == "" {
63− shellPath = "/bin/bash"
64− }
65−
66− // Default shell args
67− shellArgs := []string{"-l"}
68−
69− cmd := exec.Command(shellPath, shellArgs...)
70− cmd.Dir = cwd
71−
72− stdinPipe, err := cmd.StdinPipe()
73− if err != nil {
74− return nil
75− }
76−
77− cmd.Env = append(os.Environ(), "GIT_EDITOR=true")
78−
79− err = cmd.Start()
80− if err != nil {
81− return nil
82− }
83−
84− shell := &PersistentShell{
37+ return &PersistentShell{
85− cmd: cmd,
86− stdin: stdinPipe.(*os.File),
87− isAlive: true,
88− cwd: cwd,
38+ cwd: cwd,
89− commandQueue: make(chan *commandExecution, 10),
39+ env: os.Environ(),
90 }
40 }
91−
92− go func() {
93− defer func() {
94− if r := recover(); r != nil {
95− fmt.Fprintf(os.Stderr, "Panic in shell command processor: %v\n", r)
96− shell.isAlive = false
97− close(shell.commandQueue)
98− }
99− }()
100− shell.processCommands()
101− }()
102−
103− go func() {
104− err := cmd.Wait()
105− if err != nil {
106− // Log the error if needed
107− }
108− shell.isAlive = false
109− close(shell.commandQueue)
110− }()
111−
112− return shell
113}
41}
114
42
115−func (s *PersistentShell) processCommands() {
43+func (s *PersistentShell) Exec(ctx context.Context, command string) (string, string, error) {
116− for cmd := range s.commandQueue {
117− result := s.execCommand(cmd.command, cmd.timeout, cmd.ctx)
118− cmd.resultChan <- result
119− }
120−}
121−
122−func (s *PersistentShell) execCommand(command string, timeout time.Duration, ctx context.Context) commandResult {
123 s.mu.Lock()
44 s.mu.Lock()
124 defer s.mu.Unlock()
45 defer s.mu.Unlock()
125
46
126− if !s.isAlive {
47+ line, err := syntax.NewParser().Parse(strings.NewReader(command), "")
127− return commandResult{
128− stderr: "Shell is not alive",
129− exitCode: 1,
130− err: errors.New("shell is not alive"),
131− }
132− }
133−
134− tempDir := os.TempDir()
135− stdoutFile := filepath.Join(tempDir, fmt.Sprintf("crush-stdout-%d", time.Now().UnixNano()))
136− stderrFile := filepath.Join(tempDir, fmt.Sprintf("crush-stderr-%d", time.Now().UnixNano()))
137− statusFile := filepath.Join(tempDir, fmt.Sprintf("crush-status-%d", time.Now().UnixNano()))
138− cwdFile := filepath.Join(tempDir, fmt.Sprintf("crush-cwd-%d", time.Now().UnixNano()))
139−
140− defer func() {
141− os.Remove(stdoutFile)
142− os.Remove(stderrFile)
143− os.Remove(statusFile)
144− os.Remove(cwdFile)
145− }()
146−
147− fullCommand := fmt.Sprintf(`
148−eval %s < /dev/null > %s 2> %s
149−EXEC_EXIT_CODE=$?
150−pwd > %s
151−echo $EXEC_EXIT_CODE > %s
152−`,
153− shellQuote(command),
154− shellQuote(stdoutFile),
155− shellQuote(stderrFile),
156− shellQuote(cwdFile),
157− shellQuote(statusFile),
158− )
159−
160− _, err := s.stdin.Write([]byte(fullCommand + "\n"))
161 if err != nil {
48 if err != nil {
162− return commandResult{
49+ return "", "", fmt.Errorf("could not parse command: %w", err)
163− stderr: fmt.Sprintf("Failed to write command to shell: %v", err),
164− exitCode: 1,
165− err: err,
166− }
167 }
50 }
168
51
169− interrupted := false
52+ var stdout, stderr bytes.Buffer
170−
53+ runner, err := interp.New(
171− startTime := time.Now()
54+ interp.StdIO(nil, &stdout, &stderr),
172−
55+ interp.Interactive(false),
173− done := make(chan bool)
56+ interp.Env(expand.ListEnviron(s.env...)),
174− go func() {
57+ interp.Dir(s.cwd),
175− // Use exponential backoff polling
58+ )
176− pollInterval := 1 * time.Millisecond
177− maxPollInterval := 100 * time.Millisecond
178−
179− ticker := time.NewTicker(pollInterval)
180− defer ticker.Stop()
181−
182− for {
183− select {
184− case <-ctx.Done():
185− s.killChildren()
186− interrupted = true
187− done <- true
188− return
189−
190− case <-ticker.C:
191− if fileExists(statusFile) && fileSize(statusFile) > 0 {
192− done <- true
193− return
194− }
195−
196− if timeout > 0 {
197− elapsed := time.Since(startTime)
198− if elapsed > timeout {
199− s.killChildren()
200− interrupted = true
201− done <- true
202− return
203− }
204− }
205−
206− // Exponential backoff to reduce CPU usage for longer-running commands
207− if pollInterval < maxPollInterval {
208− pollInterval = min(time.Duration(float64(pollInterval)*1.5), maxPollInterval)
209− ticker.Reset(pollInterval)
210− }
211− }
212− }
213− }()
214−
215− <-done
216−
217− stdout := readFileOrEmpty(stdoutFile)
218− stderr := readFileOrEmpty(stderrFile)
219− exitCodeStr := readFileOrEmpty(statusFile)
220− newCwd := readFileOrEmpty(cwdFile)
221−
222− exitCode := 0
223− if exitCodeStr != "" {
224− fmt.Sscanf(exitCodeStr, "%d", &exitCode)
225− } else if interrupted {
226− exitCode = 143
227− stderr += "\nCommand execution timed out or was interrupted"
228− }
229−
230− if newCwd != "" {
231− s.cwd = strings.TrimSpace(newCwd)
232− }
233−
234− return commandResult{
235− stdout: stdout,
236− stderr: stderr,
237− exitCode: exitCode,
238− interrupted: interrupted,
239− }
240−}
241−
242−func (s *PersistentShell) killChildren() {
243− if s.cmd == nil || s.cmd.Process == nil {
244− return
245− }
246−
247− pgrepCmd := exec.Command("pgrep", "-P", fmt.Sprintf("%d", s.cmd.Process.Pid))
248− output, err := pgrepCmd.Output()
249 if err != nil {
59 if err != nil {
250− return
60+ return "", "", fmt.Errorf("could not run command: %w", err)
251 }
61 }
252
62
253− for pidStr := range strings.SplitSeq(string(output), "\n") {
63+ err = runner.Run(ctx, line)
254− if pidStr = strings.TrimSpace(pidStr); pidStr != "" {
64+ s.cwd = runner.Dir
255− var pid int
65+ s.env = []string{}
256− fmt.Sscanf(pidStr, "%d", &pid)
66+ for name, vr := range runner.Vars {
257− if pid > 0 {
67+ s.env = append(s.env, fmt.Sprintf("%s=%s", name, vr.Str))
258− proc, err := os.FindProcess(pid)
259− if err == nil {
260− proc.Signal(syscall.SIGTERM)
261− }
262− }
263− }
264 }
68 }
69+ logging.InfoPersist("Command finished", "command", command, "err", err)
70+ return stdout.String(), stderr.String(), err
265}
71}
266
72
267−func (s *PersistentShell) Exec(ctx context.Context, command string, timeoutMs int) (string, string, int, bool, error) {
73+func IsInterrupt(err error) bool {
268− if !s.isAlive {
74+ return errors.Is(err, context.Canceled) ||
269− return "", "Shell is not alive", 1, false, errors.New("shell is not alive")
75+ errors.Is(err, context.DeadlineExceeded)
270− }
271−
272− timeout := time.Duration(timeoutMs) * time.Millisecond
273−
274− resultChan := make(chan commandResult)
275− s.commandQueue <- &commandExecution{
276− command: command,
277− timeout: timeout,
278− resultChan: resultChan,
279− ctx: ctx,
280− }
281−
282− result := <-resultChan
283− return result.stdout, result.stderr, result.exitCode, result.interrupted, result.err
284}
76}
285
77
286−func (s *PersistentShell) Close() {
78+func ExitCode(err error) int {
287− s.mu.Lock()
79+ if err == nil {
288− defer s.mu.Unlock()
289−
290− if !s.isAlive {
291− return
80+ return 0
292 }
81 }
293−
82+ status, ok := interp.IsExitStatus(err)
294− s.stdin.Write([]byte("exit\n"))
83+ if ok {
295−
296− s.cmd.Process.Kill()
297− s.isAlive = false
298−}
299−
300−func shellQuote(s string) string {
301− return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
302−}
303−
304−func readFileOrEmpty(path string) string {
305− content, err := os.ReadFile(path)
306− if err != nil {
307− return ""
84+ return int(status)
308 }
85 }
309− return string(content)
86+ return 1
310−}
311−
312−func fileExists(path string) bool {
313− _, err := os.Stat(path)
314− return err == nil
315}
87}
316−
317−func fileSize(path string) int64 {
318− info, err := os.Stat(path)
319− if err != nil {
320− return 0
321− }
322− return info.Size()
323−}
internal/llm/tools/shell/shell_test.go+65−12
@@ -2,28 +2,81 @@ package shell
2
2
3import (
3import (
4 "context"
4 "context"
5− "os"
6 "testing"
5 "testing"
7−
6+ "time"
8− "github.com/stretchr/testify/require"
9)
7)
10
8
11// Benchmark to measure CPU efficiency
9// Benchmark to measure CPU efficiency
12func BenchmarkShellQuickCommands(b *testing.B) {
10func BenchmarkShellQuickCommands(b *testing.B) {
13− tmpDir, err := os.MkdirTemp("", "shell-bench")
11+ shell := newPersistentShell(b.TempDir())
14− require.NoError(b, err)
15− defer os.RemoveAll(tmpDir)
16−
17− shell := GetPersistentShell(tmpDir)
18− defer shell.Close()
19
12
20− b.ResetTimer()
21 b.ReportAllocs()
13 b.ReportAllocs()
22
14
23 for b.Loop() {
15 for b.Loop() {
24− _, _, exitCode, _, err := shell.Exec(context.Background(), "echo test", 0)
16+ _, _, err := shell.Exec(context.Background(), "echo test")
17+ exitCode := ExitCode(err)
25 if err != nil || exitCode != 0 {
18 if err != nil || exitCode != 0 {
26 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
19 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
27 }
20 }
28 }
21 }
29−}
22+}
23+
24+func TestTestTimeout(t *testing.T) {
25+ ctx, cancel := context.WithTimeout(t.Context(), time.Millisecond)
26+ t.Cleanup(cancel)
27+
28+ shell := newPersistentShell(t.TempDir())
29+ _, _, err := shell.Exec(ctx, "sleep 10")
30+ if status := ExitCode(err); status == 0 {
31+ t.Fatalf("Expected non-zero exit status, got %d", status)
32+ }
33+ if !IsInterrupt(err) {
34+ t.Fatalf("Expected command to be interrupted, but it was not")
35+ }
36+ if err == nil {
37+ t.Fatalf("Expected an error due to timeout, but got none")
38+ }
39+}
40+
41+func TestTestCancel(t *testing.T) {
42+ ctx, cancel := context.WithCancel(t.Context())
43+ cancel() // immediately cancel the context
44+
45+ shell := newPersistentShell(t.TempDir())
46+ _, _, err := shell.Exec(ctx, "sleep 10")
47+ if status := ExitCode(err); status == 0 {
48+ t.Fatalf("Expected non-zero exit status, got %d", status)
49+ }
50+ if !IsInterrupt(err) {
51+ t.Fatalf("Expected command to be interrupted, but it was not")
52+ }
53+ if err == nil {
54+ t.Fatalf("Expected an error due to cancel, but got none")
55+ }
56+}
57+
58+func TestRunCommandError(t *testing.T) {
59+ shell := newPersistentShell(t.TempDir())
60+ _, _, err := shell.Exec(t.Context(), "nopenopenope")
61+ if status := ExitCode(err); status == 0 {
62+ t.Fatalf("Expected non-zero exit status, got %d", status)
63+ }
64+ if IsInterrupt(err) {
65+ t.Fatalf("Expected command to not be interrupted, but it was")
66+ }
67+ if err == nil {
68+ t.Fatalf("Expected an error, got nil")
69+ }
70+}
71+
72+func TestRunContinuity(t *testing.T) {
73+ shell := newPersistentShell(t.TempDir())
74+ shell.Exec(t.Context(), "export FOO=bar")
75+ dst := t.TempDir()
76+ shell.Exec(t.Context(), "cd "+dst)
77+ out, _, _ := shell.Exec(t.Context(), "echo $FOO ; pwd")
78+ expect := "bar\n" + dst + "\n"
79+ if out != expect {
80+ t.Fatalf("Expected output %q, got %q", expect, out)
81+ }
82+}