1package tools
2
3import (
4 "bytes"
5 "context"
6 "crypto/rand"
7 "encoding/hex"
8 "encoding/json"
9 "fmt"
10 "io"
11 "os"
12 "os/exec"
13 "path/filepath"
14 "runtime"
15 "strings"
16 "time"
17
18 "github.com/patriceckhart/zot/packages/core"
19 "github.com/patriceckhart/zot/packages/provider"
20)
21
22const (
23 maxBashLines = 2000
24 maxBashBytes = 50 * 1024
25)
26
27// BashTool runs a shell command in the agent's cwd.
28type BashTool struct {
29 CWD string
30 Sandbox *Sandbox
31}
32
33type bashArgs struct {
34 Command string `json:"command"`
35 Timeout int `json:"timeout,omitempty"`
36}
37
38const bashSchema = `{"type":"object","properties":{"command":{"type":"string"},"timeout":{"type":"integer"}},"required":["command"]}`
39
40func (t *BashTool) Name() string { return "bash" }
41func (t *BashTool) Description() string { return shellDescription(currentShell()) }
42func (t *BashTool) Schema() json.RawMessage { return json.RawMessage(bashSchema) }
43
44func (t *BashTool) Execute(ctx context.Context, raw json.RawMessage, progress func(string)) (core.ToolResult, error) {
45 var a bashArgs
46 if err := json.Unmarshal(raw, &a); err != nil {
47 return core.ToolResult{}, fmt.Errorf("invalid args: %w", err)
48 }
49 if strings.TrimSpace(a.Command) == "" {
50 return core.ToolResult{}, fmt.Errorf("command is required")
51 }
52 if err := t.Sandbox.CheckCommand(a.Command); err != nil {
53 return core.ToolResult{}, err
54 }
55 if err := t.Sandbox.CheckBashPermission(a.Command); err != nil {
56 return core.ToolResult{}, err
57 }
58 cwd := t.CWD
59 if cwd == "" {
60 cwd, _ = os.Getwd()
61 }
62
63 runCtx := ctx
64 var cancel context.CancelFunc
65 if a.Timeout > 0 {
66 runCtx, cancel = context.WithTimeout(ctx, time.Duration(a.Timeout)*time.Second)
67 defer cancel()
68 }
69
70 start := time.Now()
71 cmd := newShellCmd(runCtx, a.Command)
72 cmd.Dir = cwd
73 cmd.Env = os.Environ()
74 setProcessGroup(cmd)
75
76 // Capture merged stdout+stderr with line-by-line streaming.
77 pr, pw := io.Pipe()
78 cmd.Stdout = pw
79 cmd.Stderr = pw
80
81 if err := cmd.Start(); err != nil {
82 return core.ToolResult{}, fmt.Errorf("start: %w", err)
83 }
84
85 // Writer to both the buffer (trimmed) and progress callback.
86 captured := &bytes.Buffer{}
87 done := make(chan struct{})
88
89 // Watch for context cancellation and kill the entire process
90 // group immediately. exec.CommandContext only kills the direct
91 // process, but child processes (e.g. grep spawned by the shell)
92 // keep the output pipe open and block cmd.Wait() indefinitely.
93 go func() {
94 select {
95 case <-runCtx.Done():
96 killProcessGroup(cmd)
97 // Close the write end so the reader goroutine unblocks.
98 pw.Close()
99 case <-done:
100 }
101 }()
102 go func() {
103 defer close(done)
104 buf := make([]byte, 4096)
105 for {
106 n, err := pr.Read(buf)
107 if n > 0 {
108 chunk := buf[:n]
109 if captured.Len() < maxBashBytes {
110 room := maxBashBytes - captured.Len()
111 if n > room {
112 captured.Write(chunk[:room])
113 } else {
114 captured.Write(chunk)
115 }
116 }
117 if progress != nil {
118 progress(string(chunk))
119 }
120 }
121 if err != nil {
122 return
123 }
124 }
125 }()
126
127 waitErr := cmd.Wait()
128 pw.Close()
129 <-done
130
131 output := captured.String()
132 truncBytes := captured.Len() >= maxBashBytes
133 lines := strings.Split(output, "\n")
134 truncLines := false
135 if len(lines) > maxBashLines {
136 lines = lines[:maxBashLines]
137 truncLines = true
138 }
139 trimmed := strings.Join(lines, "\n")
140
141 exitCode := 0
142 if waitErr != nil {
143 if ee, ok := waitErr.(*exec.ExitError); ok {
144 exitCode = ee.ExitCode()
145 } else {
146 exitCode = -1
147 }
148 }
149
150 elapsed := time.Since(start)
151
152 // Terminal-log style: echo the command on the first line with
153 // a shell-prompt prefix, a blank line, the captured output, and
154 // a footer showing exit code + elapsed time. Matches the look
155 // a human would see if they ran the command themselves, which
156 // makes the model's reasoning about it more natural too.
157 var sb strings.Builder
158 fmt.Fprintf(&sb, "$ %s\n", a.Command)
159 if trimmed != "" {
160 sb.WriteString("\n")
161 sb.WriteString(trimmed)
162 if !strings.HasSuffix(trimmed, "\n") {
163 sb.WriteString("\n")
164 }
165 }
166 if truncLines {
167 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxBashLines)
168 }
169 if truncBytes {
170 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxBashBytes)
171 }
172 sb.WriteString("\n")
173 if exitCode == 0 {
174 fmt.Fprintf(&sb, "[exit 0]")
175 } else {
176 fmt.Fprintf(&sb, "[exit %d]", exitCode)
177 }
178
179 var fullPath string
180 if truncBytes || truncLines {
181 fullPath = writeFullOutput(output)
182 if fullPath != "" {
183 fmt.Fprintf(&sb, " (full output: %s)", fullPath)
184 }
185 }
186 fmt.Fprintf(&sb, " Took %s", humanDuration(elapsed))
187
188 isErr := exitCode != 0 || ctx.Err() != nil
189 return core.ToolResult{
190 Content: []provider.Content{provider.TextBlock{Text: sb.String()}},
191 IsError: isErr,
192 Details: map[string]any{
193 "exit_code": exitCode,
194 "full_output_path": fullPath,
195 "lines_truncated": truncLines,
196 "bytes_truncated": truncBytes,
197 "duration_ms": elapsed.Milliseconds(),
198 },
199 }, nil
200}
201
202// humanDuration renders a duration in the "Took X.Ys" style used by
203// the shell-log output: tenths of a second for sub-minute runs,
204// whole seconds once we pass a minute. Trailing zeros dropped so
205// "0.1s" instead of "0.10s".
206func humanDuration(d time.Duration) string {
207 switch {
208 case d < time.Millisecond:
209 return "0.0s"
210 case d < time.Minute:
211 s := d.Seconds()
212 return fmt.Sprintf("%.1fs", s)
213 case d < time.Hour:
214 m := int(d.Minutes())
215 s := int(d.Seconds()) - m*60
216 return fmt.Sprintf("%dm%ds", m, s)
217 default:
218 h := int(d.Hours())
219 m := int(d.Minutes()) - h*60
220 return fmt.Sprintf("%dh%dm", h, m)
221 }
222}
223
224func writeFullOutput(s string) string {
225 b := make([]byte, 6)
226 _, _ = rand.Read(b)
227 name := filepath.Join(os.TempDir(), "zot-bash-"+hex.EncodeToString(b)+".log")
228 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
229 return ""
230 }
231 return name
232}
233
234type shellCommand struct {
235 path string
236 flag string
237 isBash bool
238}
239
240func currentShell() shellCommand {
241 return resolveShell(runtime.GOOS, isExecutableFile, exec.LookPath)
242}
243
244func resolveShell(goos string, executable func(string) bool, lookPath func(string) (string, error)) shellCommand {
245 if goos == "windows" {
246 return shellCommand{path: "cmd", flag: "/C"}
247 }
248 if executable("/bin/bash") {
249 return shellCommand{path: "/bin/bash", flag: "-c", isBash: true}
250 }
251 if path, err := lookPath("bash"); err == nil {
252 return shellCommand{path: path, flag: "-c", isBash: true}
253 }
254 return shellCommand{path: "/bin/sh", flag: "-c"}
255}
256
257func isExecutableFile(path string) bool {
258 info, err := os.Stat(path)
259 return err == nil && !info.IsDir() && info.Mode().Perm()&0o111 != 0
260}
261
262func shellDescription(shell shellCommand) string {
263 if shell.flag == "/C" {
264 return "Run a Windows Command Prompt command via cmd /C. stdout+stderr merged."
265 }
266 if shell.isBash {
267 return fmt.Sprintf("Run a Bash command via %s -c. stdout+stderr merged.", shell.path)
268 }
269 return "Bash is unavailable; run a POSIX sh command via /bin/sh -c. stdout+stderr merged."
270}
271
272func newShellCmd(ctx context.Context, command string) *exec.Cmd {
273 shell := currentShell()
274 return exec.CommandContext(ctx, shell.path, shell.flag, command)
275}