git.iko.soy/crush
on main help sign in
perf: concurrency improvements
2937e19Raphael Amorim
internal/llm/tools/shell/comparison_test.go+2−4
compare with main
unifiedsplit
@@ -13,8 +13,7 @@ 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, err := shell.Exec(t.Context(), "echo 'hello'")
16+ stdout, stderr, exitCode, _, err := shell.Exec(t.Context(), "echo 'hello'", 0)
17− exitCode := ExitCode(err)
18 duration := time.Since(start)
17 duration := time.Since(start)
19
18
20 require.NoError(t, err)
19 require.NoError(t, err)
@@ -33,8 +32,7 @@ func BenchmarkShellPolling(b *testing.B) {
33
32
34 for b.Loop() {
33 for b.Loop() {
35 // Use a short sleep to measure polling overhead
34 // Use a short sleep to measure polling overhead
36− _, _, err := shell.Exec(b.Context(), "sleep 0.02")
35+ _, _, exitCode, _, err := shell.Exec(b.Context(), "sleep 0.02", 0)
37− exitCode := ExitCode(err)
38 if err != nil || exitCode != 0 {
36 if err != nil || exitCode != 0 {
39 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
37 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
40 }
38 }
internal/llm/tools/shell/shell.go+281−45
@@ -1,87 +1,323 @@
1package shell
1package shell
2
2
3import (
3import (
4− "bytes"
5 "context"
4 "context"
6 "errors"
5 "errors"
7 "fmt"
6 "fmt"
8 "os"
7 "os"
8+ "os/exec"
9+ "path/filepath"
9 "strings"
10 "strings"
10 "sync"
11 "sync"
11−
12+ "syscall"
12− "github.com/charmbracelet/crush/internal/logging"
13+ "time"
13− "mvdan.cc/sh/v3/expand"
14− "mvdan.cc/sh/v3/interp"
15− "mvdan.cc/sh/v3/syntax"
16)
14)
17
15
18type PersistentShell struct {
16type PersistentShell struct {
19− env []string
17+ cmd *exec.Cmd
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
22}
38}
23
39
24var (
40var (
41+ shellInstance *PersistentShell
25− once sync.Once
42+ shellInstanceOnce sync.Once
26− shellInstance *PersistentShell
27)
43)
28
44
29−func GetPersistentShell(cwd string) *PersistentShell {
45+func GetPersistentShell(workingDir string) *PersistentShell {
30− once.Do(func() {
46+ shellInstanceOnce.Do(func() {
31− shellInstance = newPersistentShell(cwd)
47+ shellInstance = newPersistentShell(workingDir)
32 })
48 })
49+
50+ if shellInstance == nil {
51+ shellInstance = newPersistentShell(workingDir)
52+ } else if !shellInstance.isAlive {
53+ shellInstance = newPersistentShell(shellInstance.cwd)
54+ }
55+
33 return shellInstance
56 return shellInstance
34}
57}
35
58
36func newPersistentShell(cwd string) *PersistentShell {
59func 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+
37− return &PersistentShell{
84+ shell := &PersistentShell{
85+ cmd: cmd,
86+ stdin: stdinPipe.(*os.File),
87+ isAlive: true,
38− cwd: cwd,
88+ cwd: cwd,
39− env: os.Environ(),
89+ commandQueue: make(chan *commandExecution, 10),
40 }
90 }
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
41}
113}
42
114
43−func (s *PersistentShell) Exec(ctx context.Context, command string) (string, string, error) {
115+func (s *PersistentShell) processCommands() {
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 {
44 s.mu.Lock()
123 s.mu.Lock()
45 defer s.mu.Unlock()
124 defer s.mu.Unlock()
46
125
47− line, err := syntax.NewParser().Parse(strings.NewReader(command), "")
126+ if !s.isAlive {
48− if err != nil {
127+ return commandResult{
49− return "", "", fmt.Errorf("could not parse command: %w", err)
128+ stderr: "Shell is not alive",
129+ exitCode: 1,
130+ err: errors.New("shell is not alive"),
131+ }
50 }
132 }
51
133
52− var stdout, stderr bytes.Buffer
134+ tempDir := os.TempDir()
53− runner, err := interp.New(
135+ stdoutFile := filepath.Join(tempDir, fmt.Sprintf("crush-stdout-%d", time.Now().UnixNano()))
54− interp.StdIO(nil, &stdout, &stderr),
136+ stderrFile := filepath.Join(tempDir, fmt.Sprintf("crush-stderr-%d", time.Now().UnixNano()))
55− interp.Interactive(false),
137+ statusFile := filepath.Join(tempDir, fmt.Sprintf("crush-status-%d", time.Now().UnixNano()))
56− interp.Env(expand.ListEnviron(s.env...)),
138+ cwdFile := filepath.Join(tempDir, fmt.Sprintf("crush-cwd-%d", time.Now().UnixNano()))
57− interp.Dir(s.cwd),
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),
58 )
158 )
159+
160+ _, err := s.stdin.Write([]byte(fullCommand + "\n"))
59 if err != nil {
161 if err != nil {
60− return "", "", fmt.Errorf("could not run command: %w", err)
162+ return commandResult{
163+ stderr: fmt.Sprintf("Failed to write command to shell: %v", err),
164+ exitCode: 1,
165+ err: err,
166+ }
61 }
167 }
62
168
63− err = runner.Run(ctx, line)
169+ interrupted := false
64− s.cwd = runner.Dir
170+
65− s.env = []string{}
171+ startTime := time.Now()
66− for name, vr := range runner.Vars {
172+
67− s.env = append(s.env, fmt.Sprintf("%s=%s", name, vr.Str))
173+ done := make(chan bool)
174+ go func() {
175+ // Use exponential backoff polling
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,
68 }
239 }
69− logging.InfoPersist("Command finished", "command", command, "err", err)
70− return stdout.String(), stderr.String(), err
71}
240}
72
241
73−func IsInterrupt(err error) bool {
242+func (s *PersistentShell) killChildren() {
74− return errors.Is(err, context.Canceled) ||
243+ if s.cmd == nil || s.cmd.Process == nil {
75− errors.Is(err, context.DeadlineExceeded)
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 {
250+ return
251+ }
252+
253+ for pidStr := range strings.SplitSeq(string(output), "\n") {
254+ if pidStr = strings.TrimSpace(pidStr); pidStr != "" {
255+ var pid int
256+ fmt.Sscanf(pidStr, "%d", &pid)
257+ if pid > 0 {
258+ proc, err := os.FindProcess(pid)
259+ if err == nil {
260+ proc.Signal(syscall.SIGTERM)
261+ }
262+ }
263+ }
264+ }
76}
265}
77
266
78−func ExitCode(err error) int {
267+func (s *PersistentShell) Exec(ctx context.Context, command string, timeoutMs int) (string, string, int, bool, error) {
79− if err == nil {
268+ if !s.isAlive {
269+ return "", "Shell is not alive", 1, false, errors.New("shell is not alive")
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+}
285+
286+func (s *PersistentShell) Close() {
287+ s.mu.Lock()
288+ defer s.mu.Unlock()
289+
290+ if !s.isAlive {
80− return 0
291+ return
81 }
292 }
82− status, ok := interp.IsExitStatus(err)
293+
83− if ok {
294+ s.stdin.Write([]byte("exit\n"))
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 {
84− return int(status)
307+ return ""
85 }
308 }
86− return 1
309+ return string(content)
310+}
311+
312+func fileExists(path string) bool {
313+ _, err := os.Stat(path)
314+ return err == nil
87}
315}
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+12−65
@@ -2,81 +2,28 @@ package shell
2
2
3import (
3import (
4 "context"
4 "context"
5+ "os"
5 "testing"
6 "testing"
6− "time"
7+
8+ "github.com/stretchr/testify/require"
7)
9)
8
10
9// Benchmark to measure CPU efficiency
11// Benchmark to measure CPU efficiency
10func BenchmarkShellQuickCommands(b *testing.B) {
12func BenchmarkShellQuickCommands(b *testing.B) {
11− shell := newPersistentShell(b.TempDir())
13+ tmpDir, err := os.MkdirTemp("", "shell-bench")
14+ require.NoError(b, err)
15+ defer os.RemoveAll(tmpDir)
16+
17+ shell := GetPersistentShell(tmpDir)
18+ defer shell.Close()
12
19
20+ b.ResetTimer()
13 b.ReportAllocs()
21 b.ReportAllocs()
14
22
15 for b.Loop() {
23 for b.Loop() {
16− _, _, err := shell.Exec(context.Background(), "echo test")
24+ _, _, exitCode, _, err := shell.Exec(context.Background(), "echo test", 0)
17− exitCode := ExitCode(err)
18 if err != nil || exitCode != 0 {
25 if err != nil || exitCode != 0 {
19 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
26 b.Fatalf("Command failed: %v, exit code: %d", err, exitCode)
20 }
27 }
21 }
28 }
22−}
29+}
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−}
internal/permission/permission.go+31−7
@@ -1,10 +1,12 @@
1package permission
1package permission
2
2
3import (
3import (
4+ "context"
4 "errors"
5 "errors"
5 "path/filepath"
6 "path/filepath"
6 "slices"
7 "slices"
7 "sync"
8 "sync"
9+ "time"
8
10
9 "github.com/charmbracelet/crush/internal/config"
11 "github.com/charmbracelet/crush/internal/config"
10 "github.com/charmbracelet/crush/internal/pubsub"
12 "github.com/charmbracelet/crush/internal/pubsub"
@@ -44,9 +46,11 @@ type Service interface {
44type permissionService struct {
46type permissionService struct {
45 *pubsub.Broker[PermissionRequest]
47 *pubsub.Broker[PermissionRequest]
46
48
47− sessionPermissions []PermissionRequest
49+ sessionPermissions []PermissionRequest
50+ sessionPermissionsMu sync.RWMutex
48− pendingRequests sync.Map
51+ pendingRequests sync.Map
49− autoApproveSessions []string
52+ autoApproveSessions []string
53+ autoApproveSessionsMu sync.RWMutex
50}
54}
51
55
52func (s *permissionService) GrantPersistent(permission PermissionRequest) {
56func (s *permissionService) GrantPersistent(permission PermissionRequest) {
@@ -54,7 +58,10 @@ func (s *permissionService) GrantPersistent(permission PermissionRequest) {
54 if ok {
58 if ok {
55 respCh.(chan bool) <- true
59 respCh.(chan bool) <- true
56 }
60 }
61+
62+ s.sessionPermissionsMu.Lock()
57 s.sessionPermissions = append(s.sessionPermissions, permission)
63 s.sessionPermissions = append(s.sessionPermissions, permission)
64+ s.sessionPermissionsMu.Unlock()
58}
65}
59
66
60func (s *permissionService) Grant(permission PermissionRequest) {
67func (s *permissionService) Grant(permission PermissionRequest) {
@@ -72,9 +79,14 @@ func (s *permissionService) Deny(permission PermissionRequest) {
72}
79}
73
80
74func (s *permissionService) Request(opts CreatePermissionRequest) bool {
81func (s *permissionService) Request(opts CreatePermissionRequest) bool {
82+ s.autoApproveSessionsMu.RLock()
75− if slices.Contains(s.autoApproveSessions, opts.SessionID) {
83+ autoApprove := slices.Contains(s.autoApproveSessions, opts.SessionID)
84+ s.autoApproveSessionsMu.RUnlock()
85+
86+ if autoApprove {
76 return true
87 return true
77 }
88 }
89+
78 dir := filepath.Dir(opts.Path)
90 dir := filepath.Dir(opts.Path)
79 if dir == "." {
91 if dir == "." {
80 dir = config.WorkingDirectory()
92 dir = config.WorkingDirectory()
@@ -89,11 +101,14 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {
89 Params: opts.Params,
101 Params: opts.Params,
90 }
102 }
91
103
104+ s.sessionPermissionsMu.RLock()
92 for _, p := range s.sessionPermissions {
105 for _, p := range s.sessionPermissions {
93 if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {
106 if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {
107+ s.sessionPermissionsMu.RUnlock()
94 return true
108 return true
95 }
109 }
96 }
110 }
111+ s.sessionPermissionsMu.RUnlock()
97
112
98 respCh := make(chan bool, 1)
113 respCh := make(chan bool, 1)
99
114
@@ -102,13 +117,22 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {
102
117
103 s.Publish(pubsub.CreatedEvent, permission)
118 s.Publish(pubsub.CreatedEvent, permission)
104
119
105− // Wait for the response with a timeout
120+ // Wait for the response with a timeout to prevent indefinite blocking
121+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
122+ defer cancel()
123+
124+ select {
106− resp := <-respCh
125+ case resp := <-respCh:
107− return resp
126+ return resp
127+ case <-ctx.Done():
128+ return false // Timeout - deny by default
129+ }
108}
130}
109
131
110func (s *permissionService) AutoApproveSession(sessionID string) {
132func (s *permissionService) AutoApproveSession(sessionID string) {
133+ s.autoApproveSessionsMu.Lock()
111 s.autoApproveSessions = append(s.autoApproveSessions, sessionID)
134 s.autoApproveSessions = append(s.autoApproveSessions, sessionID)
135+ s.autoApproveSessionsMu.Unlock()
112}
136}
113
137
114func NewPermissionService() Service {
138func NewPermissionService() Service {
internal/pubsub/broker.go+2
@@ -111,6 +111,8 @@ func (b *Broker[T]) Publish(t EventType, payload T) {
111 select {
111 select {
112 case sub <- event:
112 case sub <- event:
113 default:
113 default:
114+ // Channel is full, subscriber is slow - skip this event
115+ // This prevents blocking the publisher
114 }
116 }
115 }
117 }
116}
118}