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 = `{
40 "type":"object",
41 "properties":{
42 "command":{"type":"string","description":"Shell command to execute."},
43 "timeout":{"type":"integer","description":"Timeout in seconds. No default timeout."}
44 },
45 "required":["command"],
46 "additionalProperties":false
47}`
48
49func (t *BashTool) Name() string { return "bash" }
50func (t *BashTool) Description() string { return "Run a shell command. stdout and stderr are merged." }
51func (t *BashTool) Schema() json.RawMessage { return json.RawMessage(bashSchema) }
52
53func (t *BashTool) Execute(ctx context.Context, raw json.RawMessage, progress func(string)) (core.ToolResult, error) {
54 var a bashArgs
55 if err := json.Unmarshal(raw, &a); err != nil {
56 return core.ToolResult{}, fmt.Errorf("invalid args: %w", err)
57 }
58 if strings.TrimSpace(a.Command) == "" {
59 return core.ToolResult{}, fmt.Errorf("command is required")
60 }
61 if err := t.Sandbox.CheckCommand(a.Command); err != nil {
62 return core.ToolResult{}, err
63 }
64 cwd := t.CWD
65 if cwd == "" {
66 cwd, _ = os.Getwd()
67 }
68
69 runCtx := ctx
70 var cancel context.CancelFunc
71 if a.Timeout > 0 {
72 runCtx, cancel = context.WithTimeout(ctx, time.Duration(a.Timeout)*time.Second)
73 defer cancel()
74 }
75
76 cmd := newShellCmd(runCtx, a.Command)
77 cmd.Dir = cwd
78 cmd.Env = os.Environ()
79
80 // Capture merged stdout+stderr with line-by-line streaming.
81 pr, pw := io.Pipe()
82 cmd.Stdout = pw
83 cmd.Stderr = pw
84
85 if err := cmd.Start(); err != nil {
86 return core.ToolResult{}, fmt.Errorf("start: %w", err)
87 }
88
89 // Writer to both the buffer (trimmed) and progress callback.
90 captured := &bytes.Buffer{}
91 done := make(chan struct{})
92 go func() {
93 defer close(done)
94 buf := make([]byte, 4096)
95 for {
96 n, err := pr.Read(buf)
97 if n > 0 {
98 chunk := buf[:n]
99 if captured.Len() < maxBashBytes {
100 room := maxBashBytes - captured.Len()
101 if n > room {
102 captured.Write(chunk[:room])
103 } else {
104 captured.Write(chunk)
105 }
106 }
107 if progress != nil {
108 progress(string(chunk))
109 }
110 }
111 if err != nil {
112 return
113 }
114 }
115 }()
116
117 waitErr := cmd.Wait()
118 pw.Close()
119 <-done
120
121 // On hard cancel, try to kill the process group.
122 if ctx.Err() != nil && cmd.Process != nil {
123 killProcessGroup(cmd)
124 }
125
126 output := captured.String()
127 truncBytes := captured.Len() >= maxBashBytes
128 lines := strings.Split(output, "\n")
129 truncLines := false
130 if len(lines) > maxBashLines {
131 lines = lines[:maxBashLines]
132 truncLines = true
133 }
134 trimmed := strings.Join(lines, "\n")
135
136 exitCode := 0
137 if waitErr != nil {
138 if ee, ok := waitErr.(*exec.ExitError); ok {
139 exitCode = ee.ExitCode()
140 } else {
141 exitCode = -1
142 }
143 }
144
145 var sb strings.Builder
146 if trimmed != "" {
147 sb.WriteString(trimmed)
148 if !strings.HasSuffix(trimmed, "\n") {
149 sb.WriteString("\n")
150 }
151 }
152 if truncLines {
153 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxBashLines)
154 }
155 if truncBytes {
156 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxBashBytes)
157 }
158 sb.WriteString(fmt.Sprintf("[exit %d]", exitCode))
159
160 var fullPath string
161 if truncBytes || truncLines {
162 fullPath = writeFullOutput(output)
163 if fullPath != "" {
164 fmt.Fprintf(&sb, " (full output: %s)", fullPath)
165 }
166 }
167
168 isErr := exitCode != 0 || ctx.Err() != nil
169 return core.ToolResult{
170 Content: []provider.Content{provider.TextBlock{Text: sb.String()}},
171 IsError: isErr,
172 Details: map[string]any{
173 "exit_code": exitCode,
174 "full_output_path": fullPath,
175 "lines_truncated": truncLines,
176 "bytes_truncated": truncBytes,
177 },
178 }, nil
179}
180
181func writeFullOutput(s string) string {
182 b := make([]byte, 6)
183 _, _ = rand.Read(b)
184 name := filepath.Join(os.TempDir(), "zot-bash-"+hex.EncodeToString(b)+".log")
185 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
186 return ""
187 }
188 return name
189}
190
191func newShellCmd(ctx context.Context, command string) *exec.Cmd {
192 if runtime.GOOS == "windows" {
193 return exec.CommandContext(ctx, "cmd", "/C", command)
194 }
195 return exec.CommandContext(ctx, "/bin/sh", "-c", command)
196}
197
198// killProcessGroup best-effort SIGTERM then SIGKILL.
199func killProcessGroup(cmd *exec.Cmd) {
200 if cmd.Process == nil {
201 return
202 }
203 _ = cmd.Process.Signal(syscall.SIGTERM)
204 time.AfterFunc(3*time.Second, func() {
205 _ = cmd.Process.Kill()
206 })
207}