8146006 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: improve shell

8146006Carlos Alexandro Becker1 year ago

internal/llm/tools/shell/shell.go+37−61

compare with main

unifiedsplit

@@ -1,6 +1,7 @@

1package shell

3import (

4+ "cmp"

4 "context"

5 "context"

5 "errors"

6 "errors"

6 "fmt"

7 "fmt"

@@ -67,16 +68,15 @@ func newPersistentShell(cwd string) *PersistentShell {

67 shellArgs = cfg.Shell.Args

68 shellArgs = cfg.Shell.Args

68 }

69 }

70− if shellPath == "" {

71− shellPath = os.Getenv("SHELL")

71+ shellPath = cmp.Or(shellPath, os.Getenv("SHELL"), "/bin/bash")

72− if shellPath == "" {

72+ if !strings.HasSuffix(shellPath, "bash") && !strings.HasSuffix(shellPath, "zsh") {

73+ logging.Warn("only bash and zsh are supported at this time", "shell", shellPath)

73− shellPath = "/bin/bash"

74+ shellPath = "/bin/bash"

74− }

75 }

77 // Default shell args

78 if len(shellArgs) == 0 {

79− shellArgs = []string{"-l"}

79+ shellArgs = []string{"--login"}

80 }

82 cmd := exec.Command(shellPath, shellArgs...)

@@ -127,12 +127,15 @@ func newPersistentShell(cwd string) *PersistentShell {

127

128func (s *PersistentShell) processCommands() {

129 for cmd := range s.commandQueue {

130− result := s.execCommand(cmd.command, cmd.timeout, cmd.ctx)

130+ cmd.resultChan <- s.execCommand(cmd.ctx, cmd.command, cmd.timeout)

131− cmd.resultChan <- result

132 }

131 }

133}

132}

134

133

134+const runBashCommandFormat = `%s </dev/null >%q 2>%q

135+echo $? >%q

136+pwd >%q`

137+

135−func (s *PersistentShell) execCommand(command string, timeout time.Duration, ctx context.Context) commandResult {

138+func (s *PersistentShell) execCommand(ctx context.Context, command string, timeout time.Duration) commandResult {

136 s.mu.Lock()

139 s.mu.Lock()

137 defer s.mu.Unlock()

140 defer s.mu.Unlock()

138

141

@@ -144,34 +147,22 @@ func (s *PersistentShell) execCommand(command string, timeout time.Duration, ctx

144 }

147 }

145 }

148 }

146

149

147− tempDir := os.TempDir()

150+ tmp := os.TempDir()

151+ now := time.Now().UnixNano()

148− stdoutFile := filepath.Join(tempDir, fmt.Sprintf("crush-stdout-%d", time.Now().UnixNano()))

152+ stdoutFile := filepath.Join(tmp, fmt.Sprintf("crush-stdout-%d", now))

149− stderrFile := filepath.Join(tempDir, fmt.Sprintf("crush-stderr-%d", time.Now().UnixNano()))

153+ stderrFile := filepath.Join(tmp, fmt.Sprintf("crush-stderr-%d", now))

150− statusFile := filepath.Join(tempDir, fmt.Sprintf("crush-status-%d", time.Now().UnixNano()))

154+ statusFile := filepath.Join(tmp, fmt.Sprintf("crush-status-%d", now))

151− cwdFile := filepath.Join(tempDir, fmt.Sprintf("crush-cwd-%d", time.Now().UnixNano()))

155+ cwdFile := filepath.Join(tmp, fmt.Sprintf("crush-cwd-%d", now))

152

156

153 defer func() {

157 defer func() {

154− os.Remove(stdoutFile)

158+ _ = os.Remove(stdoutFile)

155− os.Remove(stderrFile)

159+ _ = os.Remove(stderrFile)

156− os.Remove(statusFile)

160+ _ = os.Remove(statusFile)

157− os.Remove(cwdFile)

161+ _ = os.Remove(cwdFile)

158 }()

162 }()

159

163

160− fullCommand := fmt.Sprintf(`

164+ script := fmt.Sprintf(runBashCommandFormat, command, stdoutFile, stderrFile, statusFile, cwdFile)

161−eval %s < /dev/null > %s 2> %s

162−EXEC_EXIT_CODE=$?

163−pwd > %s

164−echo $EXEC_EXIT_CODE > %s

165−`,

166− shellQuote(command),

167− shellQuote(stdoutFile),

168− shellQuote(stderrFile),

169− shellQuote(cwdFile),

170− shellQuote(statusFile),

171− )

172−

173− _, err := s.stdin.Write([]byte(fullCommand + "\n"))

165+ if _, err := s.stdin.Write([]byte(script + "\n")); err != nil {

174− if err != nil {

175 return commandResult{

166 return commandResult{

176 stderr: fmt.Sprintf("Failed to write command to shell: %v", err),

167 stderr: fmt.Sprintf("Failed to write command to shell: %v", err),

177 exitCode: 1,

168 exitCode: 1,

@@ -180,18 +171,18 @@ echo $EXEC_EXIT_CODE > %s

180 }

171 }

181

172

182 interrupted := false

173 interrupted := false

183−

184− startTime := time.Now()

185−

186 done := make(chan bool)

174 done := make(chan bool)

187 go func() {

175 go func() {

188 // Use exponential backoff polling

176 // Use exponential backoff polling

189− pollInterval := 1 * time.Millisecond

177+ pollInterval := 10 * time.Millisecond

190− maxPollInterval := 100 * time.Millisecond

178+ maxPollInterval := time.Second

191

179

192 ticker := time.NewTicker(pollInterval)

180 ticker := time.NewTicker(pollInterval)

193 defer ticker.Stop()

181 defer ticker.Stop()

194

182

183+ timeoutTicker := time.NewTicker(cmp.Or(timeout, time.Hour*99999))

184+ defer timeoutTicker.Stop()

185+

195 for {

186 for {

196 select {

187 select {

197 case <-ctx.Done():

188 case <-ctx.Done():

@@ -200,22 +191,18 @@ echo $EXEC_EXIT_CODE > %s

200 done <- true

191 done <- true

201 return

192 return

202

193

194+ case <-timeoutTicker.C:

195+ s.killChildren()

196+ interrupted = true

197+ done <- true

198+ return

199+

203 case <-ticker.C:

200 case <-ticker.C:

204− if fileExists(statusFile) && fileSize(statusFile) > 0 {

201+ if fileSize(statusFile) > 0 {

205 done <- true

202 done <- true

206 return

203 return

207 }

204 }

208

205

209− if timeout > 0 {

210− elapsed := time.Since(startTime)

211− if elapsed > timeout {

212− s.killChildren()

213− interrupted = true

214− done <- true

215− return

216− }

217− }

218−

219 // Exponential backoff to reduce CPU usage for longer-running commands

206 // Exponential backoff to reduce CPU usage for longer-running commands

220 if pollInterval < maxPollInterval {

207 if pollInterval < maxPollInterval {

221 pollInterval = min(time.Duration(float64(pollInterval)*1.5), maxPollInterval)

208 pollInterval = min(time.Duration(float64(pollInterval)*1.5), maxPollInterval)

@@ -280,12 +267,10 @@ func (s *PersistentShell) Exec(ctx context.Context, command string, timeoutMs in

280 return "", "Shell is not alive", 1, false, errors.New("shell is not alive")

267 return "", "Shell is not alive", 1, false, errors.New("shell is not alive")

281 }

268 }

282

269

283− timeout := time.Duration(timeoutMs) * time.Millisecond

284−

285 resultChan := make(chan commandResult)

270 resultChan := make(chan commandResult)

286 s.commandQueue <- &commandExecution{

271 s.commandQueue <- &commandExecution{

287 command: command,

272 command: command,

288− timeout: timeout,

273+ timeout: time.Duration(timeoutMs) * time.Millisecond,

289 resultChan: resultChan,

274 resultChan: resultChan,

290 ctx: ctx,

275 ctx: ctx,

291 }

276 }

@@ -310,10 +295,6 @@ func (s *PersistentShell) Close() {

310 s.isAlive = false

295 s.isAlive = false

311}

296}

312

297

313−func shellQuote(s string) string {

314− return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"

315−}

316−

317func readFileOrEmpty(path string) string {

298func readFileOrEmpty(path string) string {

318 content, err := os.ReadFile(path)

299 content, err := os.ReadFile(path)

319 if err != nil {

300 if err != nil {

@@ -322,11 +303,6 @@ func readFileOrEmpty(path string) string {

322 return string(content)

303 return string(content)

323}

304}

324

305

325−func fileExists(path string) bool {

326− _, err := os.Stat(path)

327− return err == nil

328−}

329−

330func fileSize(path string) int64 {

306func fileSize(path string) int64 {

331 info, err := os.Stat(path)

307 info, err := os.Stat(path)

332 if err != nil {

308 if err != nil {