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 "syscall"
17 "time"
18
19 "github.com/patriceckhart/zot/internal/core"
20 "github.com/patriceckhart/zot/internal/provider"
21)
22
23const (
24 maxBashLines = 2000
25 maxBashBytes = 50 * 1024
26)
27
28// BashTool runs a shell command in the agent's cwd.
29type BashTool struct {
30 CWD string
31 Sandbox *Sandbox
32}
33
34type bashArgs struct {
35 Command string `json:"command"`
36 Timeout int `json:"timeout,omitempty"`
37}
38
39const bashSchema = `{"type":"object","properties":{"command":{"type":"string"},"timeout":{"type":"integer"}},"required":["command"]}`
40
41func (t *BashTool) Name() string { return "bash" }
42func (t *BashTool) Description() string { return "Run a shell command. stdout+stderr merged." }
43func (t *BashTool) Schema() json.RawMessage { return json.RawMessage(bashSchema) }
44
45func (t *BashTool) Execute(ctx context.Context, raw json.RawMessage, progress func(string)) (core.ToolResult, error) {
46 var a bashArgs
47 if err := json.Unmarshal(raw, &a); err != nil {
48 return core.ToolResult{}, fmt.Errorf("invalid args: %w", err)
49 }
50 if strings.TrimSpace(a.Command) == "" {
51 return core.ToolResult{}, fmt.Errorf("command is required")
52 }
53 if err := t.Sandbox.CheckCommand(a.Command); err != nil {
54 return core.ToolResult{}, err
55 }
56 cwd := t.CWD
57 if cwd == "" {
58 cwd, _ = os.Getwd()
59 }
60
61 runCtx := ctx
62 var cancel context.CancelFunc
63 if a.Timeout > 0 {
64 runCtx, cancel = context.WithTimeout(ctx, time.Duration(a.Timeout)*time.Second)
65 defer cancel()
66 }
67
68 start := time.Now()
69 cmd := newShellCmd(runCtx, a.Command)
70 cmd.Dir = cwd
71 cmd.Env = os.Environ()
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 go func() {
86 defer close(done)
87 buf := make([]byte, 4096)
88 for {
89 n, err := pr.Read(buf)
90 if n > 0 {
91 chunk := buf[:n]
92 if captured.Len() < maxBashBytes {
93 room := maxBashBytes - captured.Len()
94 if n > room {
95 captured.Write(chunk[:room])
96 } else {
97 captured.Write(chunk)
98 }
99 }
100 if progress != nil {
101 progress(string(chunk))
102 }
103 }
104 if err != nil {
105 return
106 }
107 }
108 }()
109
110 waitErr := cmd.Wait()
111 pw.Close()
112 <-done
113
114 // On hard cancel, try to kill the process group.
115 if ctx.Err() != nil && cmd.Process != nil {
116 killProcessGroup(cmd)
117 }
118
119 output := captured.String()
120 truncBytes := captured.Len() >= maxBashBytes
121 lines := strings.Split(output, "\n")
122 truncLines := false
123 if len(lines) > maxBashLines {
124 lines = lines[:maxBashLines]
125 truncLines = true
126 }
127 trimmed := strings.Join(lines, "\n")
128
129 exitCode := 0
130 if waitErr != nil {
131 if ee, ok := waitErr.(*exec.ExitError); ok {
132 exitCode = ee.ExitCode()
133 } else {
134 exitCode = -1
135 }
136 }
137
138 elapsed := time.Since(start)
139
140 // Terminal-log style: echo the command on the first line with
141 // a shell-prompt prefix, a blank line, the captured output, and
142 // a footer showing exit code + elapsed time. Matches the look
143 // a human would see if they ran the command themselves, which
144 // makes the model's reasoning about it more natural too.
145 var sb strings.Builder
146 fmt.Fprintf(&sb, "$ %s\n", a.Command)
147 if trimmed != "" {
148 sb.WriteString("\n")
149 sb.WriteString(trimmed)
150 if !strings.HasSuffix(trimmed, "\n") {
151 sb.WriteString("\n")
152 }
153 }
154 if truncLines {
155 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxBashLines)
156 }
157 if truncBytes {
158 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxBashBytes)
159 }
160 sb.WriteString("\n")
161 if exitCode == 0 {
162 fmt.Fprintf(&sb, "[exit 0]")
163 } else {
164 fmt.Fprintf(&sb, "[exit %d]", exitCode)
165 }
166
167 var fullPath string
168 if truncBytes || truncLines {
169 fullPath = writeFullOutput(output)
170 if fullPath != "" {
171 fmt.Fprintf(&sb, " (full output: %s)", fullPath)
172 }
173 }
174 fmt.Fprintf(&sb, " Took %s", humanDuration(elapsed))
175
176 isErr := exitCode != 0 || ctx.Err() != nil
177 return core.ToolResult{
178 Content: []provider.Content{provider.TextBlock{Text: sb.String()}},
179 IsError: isErr,
180 Details: map[string]any{
181 "exit_code": exitCode,
182 "full_output_path": fullPath,
183 "lines_truncated": truncLines,
184 "bytes_truncated": truncBytes,
185 "duration_ms": elapsed.Milliseconds(),
186 },
187 }, nil
188}
189
190// humanDuration renders a duration in the "Took X.Ys" style used by
191// the shell-log output: tenths of a second for sub-minute runs,
192// whole seconds once we pass a minute. Trailing zeros dropped so
193// "0.1s" instead of "0.10s".
194func humanDuration(d time.Duration) string {
195 switch {
196 case d < time.Millisecond:
197 return "0.0s"
198 case d < time.Minute:
199 s := d.Seconds()
200 return fmt.Sprintf("%.1fs", s)
201 case d < time.Hour:
202 m := int(d.Minutes())
203 s := int(d.Seconds()) - m*60
204 return fmt.Sprintf("%dm%ds", m, s)
205 default:
206 h := int(d.Hours())
207 m := int(d.Minutes()) - h*60
208 return fmt.Sprintf("%dh%dm", h, m)
209 }
210}
211
212func writeFullOutput(s string) string {
213 b := make([]byte, 6)
214 _, _ = rand.Read(b)
215 name := filepath.Join(os.TempDir(), "zot-bash-"+hex.EncodeToString(b)+".log")
216 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
217 return ""
218 }
219 return name
220}
221
222func newShellCmd(ctx context.Context, command string) *exec.Cmd {
223 if runtime.GOOS == "windows" {
224 return exec.CommandContext(ctx, "cmd", "/C", command)
225 }
226 return exec.CommandContext(ctx, "/bin/sh", "-c", command)
227}
228
229// killProcessGroup best-effort SIGTERM then SIGKILL.
230func killProcessGroup(cmd *exec.Cmd) {
231 if cmd.Process == nil {
232 return
233 }
234 _ = cmd.Process.Signal(syscall.SIGTERM)
235 time.AfterFunc(3*time.Second, func() {
236 _ = cmd.Process.Kill()
237 })
238}