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