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/internal/core"
19 "github.com/patriceckhart/zot/internal/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 "Run a shell command. stdout+stderr merged." }
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 cwd := t.CWD
56 if cwd == "" {
57 cwd, _ = os.Getwd()
58 }
59
60 runCtx := ctx
61 var cancel context.CancelFunc
62 if a.Timeout > 0 {
63 runCtx, cancel = context.WithTimeout(ctx, time.Duration(a.Timeout)*time.Second)
64 defer cancel()
65 }
66
67 start := time.Now()
68 cmd := newShellCmd(runCtx, a.Command)
69 cmd.Dir = cwd
70 cmd.Env = os.Environ()
71 setProcessGroup(cmd)
72
73 // Capture merged stdout+stderr with line-by-line streaming.
74 pr, pw := io.Pipe()
75 cmd.Stdout = pw
76 cmd.Stderr = pw
77
78 if err := cmd.Start(); err != nil {
79 return core.ToolResult{}, fmt.Errorf("start: %w", err)
80 }
81
82 // Writer to both the buffer (trimmed) and progress callback.
83 captured := &bytes.Buffer{}
84 done := make(chan struct{})
85
86 // Watch for context cancellation and kill the entire process
87 // group immediately. exec.CommandContext only kills the direct
88 // process, but child processes (e.g. grep spawned by the shell)
89 // keep the output pipe open and block cmd.Wait() indefinitely.
90 go func() {
91 select {
92 case <-runCtx.Done():
93 killProcessGroup(cmd)
94 // Close the write end so the reader goroutine unblocks.
95 pw.Close()
96 case <-done:
97 }
98 }()
99 go func() {
100 defer close(done)
101 buf := make([]byte, 4096)
102 for {
103 n, err := pr.Read(buf)
104 if n > 0 {
105 chunk := buf[:n]
106 if captured.Len() < maxBashBytes {
107 room := maxBashBytes - captured.Len()
108 if n > room {
109 captured.Write(chunk[:room])
110 } else {
111 captured.Write(chunk)
112 }
113 }
114 if progress != nil {
115 progress(string(chunk))
116 }
117 }
118 if err != nil {
119 return
120 }
121 }
122 }()
123
124 waitErr := cmd.Wait()
125 pw.Close()
126 <-done
127
128 output := captured.String()
129 truncBytes := captured.Len() >= maxBashBytes
130 lines := strings.Split(output, "\n")
131 truncLines := false
132 if len(lines) > maxBashLines {
133 lines = lines[:maxBashLines]
134 truncLines = true
135 }
136 trimmed := strings.Join(lines, "\n")
137
138 exitCode := 0
139 if waitErr != nil {
140 if ee, ok := waitErr.(*exec.ExitError); ok {
141 exitCode = ee.ExitCode()
142 } else {
143 exitCode = -1
144 }
145 }
146
147 elapsed := time.Since(start)
148
149 // Terminal-log style: echo the command on the first line with
150 // a shell-prompt prefix, a blank line, the captured output, and
151 // a footer showing exit code + elapsed time. Matches the look
152 // a human would see if they ran the command themselves, which
153 // makes the model's reasoning about it more natural too.
154 var sb strings.Builder
155 fmt.Fprintf(&sb, "$ %s\n", a.Command)
156 if trimmed != "" {
157 sb.WriteString("\n")
158 sb.WriteString(trimmed)
159 if !strings.HasSuffix(trimmed, "\n") {
160 sb.WriteString("\n")
161 }
162 }
163 if truncLines {
164 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxBashLines)
165 }
166 if truncBytes {
167 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxBashBytes)
168 }
169 sb.WriteString("\n")
170 if exitCode == 0 {
171 fmt.Fprintf(&sb, "[exit 0]")
172 } else {
173 fmt.Fprintf(&sb, "[exit %d]", exitCode)
174 }
175
176 var fullPath string
177 if truncBytes || truncLines {
178 fullPath = writeFullOutput(output)
179 if fullPath != "" {
180 fmt.Fprintf(&sb, " (full output: %s)", fullPath)
181 }
182 }
183 fmt.Fprintf(&sb, " Took %s", humanDuration(elapsed))
184
185 isErr := exitCode != 0 || ctx.Err() != nil
186 return core.ToolResult{
187 Content: []provider.Content{provider.TextBlock{Text: sb.String()}},
188 IsError: isErr,
189 Details: map[string]any{
190 "exit_code": exitCode,
191 "full_output_path": fullPath,
192 "lines_truncated": truncLines,
193 "bytes_truncated": truncBytes,
194 "duration_ms": elapsed.Milliseconds(),
195 },
196 }, nil
197}
198
199// humanDuration renders a duration in the "Took X.Ys" style used by
200// the shell-log output: tenths of a second for sub-minute runs,
201// whole seconds once we pass a minute. Trailing zeros dropped so
202// "0.1s" instead of "0.10s".
203func humanDuration(d time.Duration) string {
204 switch {
205 case d < time.Millisecond:
206 return "0.0s"
207 case d < time.Minute:
208 s := d.Seconds()
209 return fmt.Sprintf("%.1fs", s)
210 case d < time.Hour:
211 m := int(d.Minutes())
212 s := int(d.Seconds()) - m*60
213 return fmt.Sprintf("%dm%ds", m, s)
214 default:
215 h := int(d.Hours())
216 m := int(d.Minutes()) - h*60
217 return fmt.Sprintf("%dh%dm", h, m)
218 }
219}
220
221func writeFullOutput(s string) string {
222 b := make([]byte, 6)
223 _, _ = rand.Read(b)
224 name := filepath.Join(os.TempDir(), "zot-bash-"+hex.EncodeToString(b)+".log")
225 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
226 return ""
227 }
228 return name
229}
230
231func newShellCmd(ctx context.Context, command string) *exec.Cmd {
232 if runtime.GOOS == "windows" {
233 return exec.CommandContext(ctx, "cmd", "/C", command)
234 }
235 return exec.CommandContext(ctx, "/bin/sh", "-c", command)
236}