1package tools
2
3import (
4 "bytes"
5 "context"
6 "crypto/rand"
7 "encoding/hex"
8 "encoding/json"
9 "errors"
10 "fmt"
11 "io"
12 "os"
13 "os/exec"
14 "path/filepath"
15 "strings"
16 "sync"
17 "time"
18
19 "github.com/bnema/zut/packages/core"
20 "github.com/bnema/zut/packages/provider"
21)
22
23const (
24 maxPythonLines = 2000
25 maxPythonBytes = 50 * 1024
26
27 // maxPythonTimeoutSeconds is the largest whole-second timeout that can
28 // be converted to a time.Duration without overflowing it.
29 maxPythonTimeoutSeconds int64 = int64(time.Duration(1<<63-1) / time.Second)
30)
31
32// PythonTool executes inline Python 3 code against an already prepared
33// environment. Zut never installs Python, packages, uv, or virtual
34// environments. Each call runs a fresh, stateless subprocess.
35type PythonTool struct {
36 CWD string
37 Sandbox *Sandbox
38 Interpreter string
39
40 // resolveHook substitutes interpreter resolution in tests. Nil uses
41 // the default lazy resolver, which shares the caller's timeout budget.
42 resolveHook func(ctx context.Context, cwd, configured string) (resolvedPython, error)
43 // execHook substitutes process execution in tests. Nil runs a real
44 // subprocess directly (no shell, no temp file).
45 execHook func(ctx context.Context, exe, code, cwd string, progress func(string)) (pythonExecOutcome, error)
46}
47
48// pythonCancelError carries the process-termination error from the
49// platform Cancel callback to the Wait site. Cancel runs on exec's
50// internal goroutine while Wait blocks, hence the mutex.
51type pythonCancelError struct {
52 mu sync.Mutex
53 err error
54}
55
56func (e *pythonCancelError) set(err error) {
57 if e == nil {
58 return
59 }
60 e.mu.Lock()
61 defer e.mu.Unlock()
62 e.err = err
63}
64
65func (e *pythonCancelError) get() error {
66 if e == nil {
67 return nil
68 }
69 e.mu.Lock()
70 defer e.mu.Unlock()
71 return e.err
72}
73
74// pythonExecOutcome is the raw result of one Python subprocess run.
75type pythonExecOutcome struct {
76 Captured string
77 ExitCode int
78 TimedOut bool
79 Cancelled bool
80 BytesTrunc bool
81 LinesTrunc bool
82 CleanupErr error
83}
84
85type pythonArgs struct {
86 Code string `json:"code"`
87 Timeout *int64 `json:"timeout"`
88}
89
90var pythonSchema = fmt.Sprintf(`{"type":"object","properties":{"code":{"type":"string","description":"Inline Python 3 source to execute via stdin (no interactive input; stdin observes EOF). Save reusable scripts with write/edit instead."},"timeout":{"type":"integer","minimum":1,"maximum":%d,"description":"Maximum runtime in seconds, including interpreter discovery."}},"required":["code","timeout"]}`, maxPythonTimeoutSeconds)
91
92func (t *PythonTool) Name() string { return "python" }
93
94func (t *PythonTool) Description() string {
95 return "Execute inline Python 3 code with an existing interpreter (stateless stdin execution in the session cwd, merged stdout/stderr). Never installs packages; save reusable scripts with write/edit."
96}
97
98func (t *PythonTool) Schema() json.RawMessage { return json.RawMessage(pythonSchema) }
99
100func pythonTimeoutDuration(seconds int64) (time.Duration, error) {
101 if seconds < 1 || seconds > maxPythonTimeoutSeconds {
102 return 0, fmt.Errorf("timeout must be between 1 and %d seconds", maxPythonTimeoutSeconds)
103 }
104 return time.Duration(seconds) * time.Second, nil
105}
106
107// pythonExecArgv builds the direct execution argv. The executable and its
108// arguments stay separate; they are never concatenated into a shell command.
109func pythonExecArgv(exe string) []string {
110 return []string{exe, "-u", "-B", "-"}
111}
112
113func (t *PythonTool) Execute(ctx context.Context, raw json.RawMessage, progress func(string)) (core.ToolResult, error) {
114 var a pythonArgs
115 if err := json.Unmarshal(raw, &a); err != nil {
116 return protocolToolError(fmt.Sprintf("invalid args: %v", err))
117 }
118 if strings.TrimSpace(a.Code) == "" {
119 return protocolToolError("code is required")
120 }
121 if a.Timeout == nil {
122 return protocolToolError("timeout is required")
123 }
124 timeout, err := pythonTimeoutDuration(*a.Timeout)
125 if err != nil {
126 return protocolToolError(err.Error())
127 }
128 // Admission checks run before any interpreter discovery, so a denied
129 // call never probes the filesystem or spawns a process. The denial
130 // is typed (not a plain tool-error result) so core records the
131 // refusal and does not append a recovery prompt urging more actions.
132 if t.Sandbox != nil && t.Sandbox.Permissions != nil {
133 return core.ToolResult{}, denied("permission denied: this agent cannot run python")
134 }
135 if t.Sandbox.Locked() {
136 return core.ToolResult{}, denied("jailed: arbitrary Python cannot be confined by the current jail (use /unjail to disable)")
137 }
138 cwd := t.CWD
139 if cwd == "" {
140 cwd, _ = os.Getwd()
141 }
142
143 // Resolution and execution share the caller's timeout budget.
144 runCtx, cancel := context.WithTimeout(ctx, timeout)
145 defer cancel()
146
147 interp, err := t.resolve(runCtx, cwd)
148 if err != nil {
149 return toolErr(err.Error()), nil
150 }
151
152 outcome, err := t.run(runCtx, ctx, interp.Path, a.Code, cwd, progress)
153 if err != nil {
154 return core.ToolResult{}, err
155 }
156
157 trimmed, truncLines := truncatePythonLines(outcome.Captured)
158 outcome.LinesTrunc = truncLines
159 text := formatPythonResult(interp, outcome, trimmed, int64(timeout/time.Second))
160 isErr := outcome.ExitCode != 0 || outcome.TimedOut || outcome.Cancelled
161 // Like bash, persist the byte-capped buffer for truncated results so
162 // the model can read beyond the line-trimmed projection.
163 fullPath := ""
164 if outcome.BytesTrunc || outcome.LinesTrunc {
165 fullPath = writePythonFullOutput(outcome.Captured)
166 if fullPath != "" {
167 text += fmt.Sprintf(" (full output: %s)", fullPath)
168 }
169 }
170 details := map[string]any{
171 "exit_code": outcome.ExitCode,
172 "interpreter": interp.Path,
173 "python_version": interp.Version.String(),
174 "bytes_truncated": outcome.BytesTrunc,
175 "lines_truncated": outcome.LinesTrunc,
176 "full_output_path": fullPath,
177 }
178 if outcome.CleanupErr != nil {
179 details["cleanup_error"] = outcome.CleanupErr.Error()
180 }
181 return core.ToolResult{
182 Content: []provider.Content{provider.TextBlock{Text: text}},
183 IsError: isErr,
184 Details: details,
185 }, nil
186}
187
188func (t *PythonTool) resolve(ctx context.Context, cwd string) (resolvedPython, error) {
189 if t.resolveHook != nil {
190 return t.resolveHook(ctx, cwd, t.Interpreter)
191 }
192 return resolvePythonInterpreter(ctx, pythonResolveDeps{CWD: cwd, Configured: t.Interpreter})
193}
194
195func (t *PythonTool) run(runCtx, parentCtx context.Context, exe, code, cwd string, progress func(string)) (pythonExecOutcome, error) {
196 if t.execHook != nil {
197 return t.execHook(runCtx, exe, code, cwd, progress)
198 }
199 return runPythonProcess(runCtx, parentCtx, exe, code, cwd, progress)
200}
201
202// runPythonProcess runs one fresh subprocess directly with code on stdin.
203// It bounds captured output, drains pipes, and reports cleanup failure
204// without claiming every detached descendant was terminated.
205func runPythonProcess(runCtx, parentCtx context.Context, exe, code, cwd string, progress func(string)) (pythonExecOutcome, error) {
206 var outcome pythonExecOutcome
207 cmd := exec.CommandContext(runCtx, exe, "-u", "-B", "-")
208 cmd.Dir = cwd
209 cmd.Env = pythonChildEnv()
210 cmd.Stdin = strings.NewReader(code)
211
212 pr, pw := io.Pipe()
213 cmd.Stdout = pw
214 cmd.Stderr = pw
215 // The Cancel callback runs on exec's internal goroutine while Wait
216 // blocks here: capture its termination error under a mutex so a
217 // failed group/tree kill is reported as cleanup metadata instead of
218 // being folded into Wait's process error and lost.
219 var cancelErr pythonCancelError
220 closeOutput := configurePythonProcess(cmd, pw, &cancelErr)
221 defer closeOutput()
222
223 if err := cmd.Start(); err != nil {
224 _ = pw.Close()
225 return outcome, fmt.Errorf("start: %w", err)
226 }
227
228 captured := &bytes.Buffer{}
229 discarded := false
230 readerDone := make(chan struct{})
231 go func() {
232 defer close(readerDone)
233 buf := make([]byte, 4096)
234 for {
235 n, err := pr.Read(buf)
236 if n > 0 {
237 chunk := buf[:n]
238 if appendPythonChunk(captured, chunk) {
239 discarded = true
240 }
241 if progress != nil {
242 progress(string(chunk))
243 }
244 }
245 if err != nil {
246 return
247 }
248 }
249 }()
250
251 waitErr := cmd.Wait()
252 closeOutput()
253 <-readerDone
254
255 outcome.Captured = captured.String()
256 outcome.BytesTrunc = discarded
257
258 exitCode := 0
259 if waitErr != nil {
260 if ee, ok := waitErr.(*exec.ExitError); ok {
261 exitCode = ee.ExitCode()
262 } else if errors.Is(waitErr, exec.ErrWaitDelay) {
263 // The interpreter exited successfully but a detached
264 // descendant holds the output pipe: keep exit 0 and report
265 // the late cleanup as metadata instead of a failure.
266 outcome.CleanupErr = waitErr
267 } else {
268 exitCode = -1
269 }
270 }
271 outcome.ExitCode = exitCode
272 if err := cancelErr.get(); err != nil {
273 // Termination itself failed: descendants may still be running.
274 // Report it as cleanup metadata, never as an execution error
275 // when the interpreter already exited successfully.
276 if outcome.CleanupErr == nil {
277 outcome.CleanupErr = err
278 }
279 }
280 if errors.Is(runCtx.Err(), context.DeadlineExceeded) && parentCtx.Err() == nil {
281 outcome.TimedOut = true
282 } else if parentCtx.Err() != nil || errors.Is(runCtx.Err(), context.Canceled) {
283 outcome.Cancelled = true
284 }
285 return outcome, nil
286}
287
288// appendPythonChunk appends chunk to captured up to maxPythonBytes. It
289// reports whether any bytes were discarded, so output exactly filling the
290// budget is not misreported as truncated.
291func appendPythonChunk(captured *bytes.Buffer, chunk []byte) (discarded bool) {
292 if len(chunk) == 0 {
293 return false
294 }
295 if captured.Len() >= maxPythonBytes {
296 return true
297 }
298 if room := maxPythonBytes - captured.Len(); len(chunk) > room {
299 captured.Write(chunk[:room])
300 return true
301 }
302 captured.Write(chunk)
303 return false
304}
305
306// writePythonFullOutput persists the byte-capped buffer for truncated
307// results, mirroring the bash full-output artifact (zut-python prefix).
308func writePythonFullOutput(s string) string {
309 b := make([]byte, 6)
310 if _, err := rand.Read(b); err != nil {
311 return ""
312 }
313 name := filepath.Join(os.TempDir(), "zut-python-"+hex.EncodeToString(b)+".log")
314 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
315 return ""
316 }
317 return name
318}
319
320// truncatePythonLines bounds the logical line count of captured output.
321// A trailing newline does not create an extra line, matching grepOutput.
322func truncatePythonLines(output string) (trimmed string, truncLines bool) {
323 lines := strings.Split(output, "\n")
324 if strings.HasSuffix(output, "\n") && len(lines) > 0 {
325 lines = lines[:len(lines)-1]
326 }
327 if len(lines) > maxPythonLines {
328 return strings.Join(lines[:maxPythonLines], "\n"), true
329 }
330 return output, false
331}
332
333// formatPythonResult renders interpreter/version, captured output, and
334// exit/timeout/cancellation status. Truncation is marked explicitly and
335// the caller persists the full output artifact. The caller supplies the
336// line-trimmed output so Details and text agree on lines_truncated.
337func formatPythonResult(interp resolvedPython, outcome pythonExecOutcome, trimmed string, timeoutSeconds int64) string {
338 var sb strings.Builder
339 fmt.Fprintf(&sb, "python %s (%s)\n", interp.Version.String(), interp.Path)
340 if trimmed != "" {
341 sb.WriteString("\n")
342 sb.WriteString(trimmed)
343 if !strings.HasSuffix(trimmed, "\n") {
344 sb.WriteString("\n")
345 }
346 }
347 if outcome.LinesTrunc {
348 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxPythonLines)
349 }
350 if outcome.BytesTrunc {
351 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxPythonBytes)
352 }
353 sb.WriteString("\n")
354 switch {
355 case outcome.TimedOut:
356 fmt.Fprintf(&sb, "[timed out after %d second", timeoutSeconds)
357 if timeoutSeconds != 1 {
358 sb.WriteByte('s')
359 }
360 sb.WriteByte(']')
361 case outcome.Cancelled:
362 sb.WriteString("[cancelled]")
363 case outcome.ExitCode == 0:
364 sb.WriteString("[exit 0]")
365 default:
366 fmt.Fprintf(&sb, "[exit %d]", outcome.ExitCode)
367 }
368 if outcome.CleanupErr != nil {
369 fmt.Fprintf(&sb, " [cleanup: %s]", outcome.CleanupErr.Error())
370 }
371 return sb.String()
372}