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