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