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.
130 if t.Sandbox != nil && t.Sandbox.Permissions != nil {
131 return protocolToolError("permission denied: this agent cannot run python")
132 }
133 if t.Sandbox.Locked() {
134 return protocolToolError("jailed: arbitrary Python cannot be confined by the current jail (use /unjail to disable)")
135 }
136 cwd := t.CWD
137 if cwd == "" {
138 cwd, _ = os.Getwd()
139 }
140
141 // Resolution and execution share the caller's timeout budget.
142 runCtx, cancel := context.WithTimeout(ctx, timeout)
143 defer cancel()
144
145 interp, err := t.resolve(runCtx, cwd)
146 if err != nil {
147 return toolErr(err.Error()), nil
148 }
149
150 outcome, err := t.run(runCtx, ctx, interp.Path, a.Code, cwd, progress)
151 if err != nil {
152 return core.ToolResult{}, err
153 }
154
155 trimmed, truncLines := truncatePythonLines(outcome.Captured)
156 outcome.LinesTrunc = truncLines
157 text := formatPythonResult(interp, outcome, trimmed, int64(timeout/time.Second))
158 isErr := outcome.ExitCode != 0 || outcome.TimedOut || outcome.Cancelled
159 // Like bash, persist the byte-capped buffer for truncated results so
160 // the model can read beyond the line-trimmed projection.
161 fullPath := ""
162 if outcome.BytesTrunc || outcome.LinesTrunc {
163 fullPath = writePythonFullOutput(outcome.Captured)
164 if fullPath != "" {
165 text += fmt.Sprintf(" (full output: %s)", fullPath)
166 }
167 }
168 details := map[string]any{
169 "exit_code": outcome.ExitCode,
170 "interpreter": interp.Path,
171 "python_version": interp.Version.String(),
172 "bytes_truncated": outcome.BytesTrunc,
173 "lines_truncated": outcome.LinesTrunc,
174 "full_output_path": fullPath,
175 }
176 if outcome.CleanupErr != nil {
177 details["cleanup_error"] = outcome.CleanupErr.Error()
178 }
179 return core.ToolResult{
180 Content: []provider.Content{provider.TextBlock{Text: text}},
181 IsError: isErr,
182 Details: details,
183 }, nil
184}
185
186func (t *PythonTool) resolve(ctx context.Context, cwd string) (resolvedPython, error) {
187 if t.resolveHook != nil {
188 return t.resolveHook(ctx, cwd, t.Interpreter)
189 }
190 return resolvePythonInterpreter(ctx, pythonResolveDeps{CWD: cwd, Configured: t.Interpreter})
191}
192
193func (t *PythonTool) run(runCtx, parentCtx context.Context, exe, code, cwd string, progress func(string)) (pythonExecOutcome, error) {
194 if t.execHook != nil {
195 return t.execHook(runCtx, exe, code, cwd, progress)
196 }
197 return runPythonProcess(runCtx, parentCtx, exe, code, cwd, progress)
198}
199
200// runPythonProcess runs one fresh subprocess directly with code on stdin.
201// It bounds captured output, drains pipes, and reports cleanup failure
202// without claiming every detached descendant was terminated.
203func runPythonProcess(runCtx, parentCtx context.Context, exe, code, cwd string, progress func(string)) (pythonExecOutcome, error) {
204 var outcome pythonExecOutcome
205 cmd := exec.CommandContext(runCtx, exe, "-u", "-B", "-")
206 cmd.Dir = cwd
207 cmd.Env = pythonChildEnv()
208 cmd.Stdin = strings.NewReader(code)
209
210 pr, pw := io.Pipe()
211 cmd.Stdout = pw
212 cmd.Stderr = pw
213 // The Cancel callback runs on exec's internal goroutine while Wait
214 // blocks here: capture its termination error under a mutex so a
215 // failed group/tree kill is reported as cleanup metadata instead of
216 // being folded into Wait's process error and lost.
217 var cancelErr pythonCancelError
218 closeOutput := configurePythonProcess(cmd, pw, &cancelErr)
219 defer closeOutput()
220
221 if err := cmd.Start(); err != nil {
222 _ = pw.Close()
223 return outcome, fmt.Errorf("start: %w", err)
224 }
225
226 captured := &bytes.Buffer{}
227 discarded := false
228 readerDone := make(chan struct{})
229 go func() {
230 defer close(readerDone)
231 buf := make([]byte, 4096)
232 for {
233 n, err := pr.Read(buf)
234 if n > 0 {
235 chunk := buf[:n]
236 if appendPythonChunk(captured, chunk) {
237 discarded = true
238 }
239 if progress != nil {
240 progress(string(chunk))
241 }
242 }
243 if err != nil {
244 return
245 }
246 }
247 }()
248
249 waitErr := cmd.Wait()
250 closeOutput()
251 <-readerDone
252
253 outcome.Captured = captured.String()
254 outcome.BytesTrunc = discarded
255
256 exitCode := 0
257 if waitErr != nil {
258 if ee, ok := waitErr.(*exec.ExitError); ok {
259 exitCode = ee.ExitCode()
260 } else if errors.Is(waitErr, exec.ErrWaitDelay) {
261 // The interpreter exited successfully but a detached
262 // descendant holds the output pipe: keep exit 0 and report
263 // the late cleanup as metadata instead of a failure.
264 outcome.CleanupErr = waitErr
265 } else {
266 exitCode = -1
267 }
268 }
269 outcome.ExitCode = exitCode
270 if err := cancelErr.get(); err != nil {
271 // Termination itself failed: descendants may still be running.
272 // Report it as cleanup metadata, never as an execution error
273 // when the interpreter already exited successfully.
274 if outcome.CleanupErr == nil {
275 outcome.CleanupErr = err
276 }
277 }
278 if errors.Is(runCtx.Err(), context.DeadlineExceeded) && parentCtx.Err() == nil {
279 outcome.TimedOut = true
280 } else if parentCtx.Err() != nil || errors.Is(runCtx.Err(), context.Canceled) {
281 outcome.Cancelled = true
282 }
283 return outcome, nil
284}
285
286// appendPythonChunk appends chunk to captured up to maxPythonBytes. It
287// reports whether any bytes were discarded, so output exactly filling the
288// budget is not misreported as truncated.
289func appendPythonChunk(captured *bytes.Buffer, chunk []byte) (discarded bool) {
290 if len(chunk) == 0 {
291 return false
292 }
293 if captured.Len() >= maxPythonBytes {
294 return true
295 }
296 if room := maxPythonBytes - captured.Len(); len(chunk) > room {
297 captured.Write(chunk[:room])
298 return true
299 }
300 captured.Write(chunk)
301 return false
302}
303
304// writePythonFullOutput persists the byte-capped buffer for truncated
305// results, mirroring the bash full-output artifact (zut-python prefix).
306func writePythonFullOutput(s string) string {
307 b := make([]byte, 6)
308 if _, err := rand.Read(b); err != nil {
309 return ""
310 }
311 name := filepath.Join(os.TempDir(), "zut-python-"+hex.EncodeToString(b)+".log")
312 if err := os.WriteFile(name, []byte(s), 0o600); err != nil {
313 return ""
314 }
315 return name
316}
317
318// truncatePythonLines bounds the logical line count of captured output.
319// A trailing newline does not create an extra line, matching grepOutput.
320func truncatePythonLines(output string) (trimmed string, truncLines bool) {
321 lines := strings.Split(output, "\n")
322 if strings.HasSuffix(output, "\n") && len(lines) > 0 {
323 lines = lines[:len(lines)-1]
324 }
325 if len(lines) > maxPythonLines {
326 return strings.Join(lines[:maxPythonLines], "\n"), true
327 }
328 return output, false
329}
330
331// formatPythonResult renders interpreter/version, captured output, and
332// exit/timeout/cancellation status. Truncation is marked explicitly and
333// the caller persists the full output artifact. The caller supplies the
334// line-trimmed output so Details and text agree on lines_truncated.
335func formatPythonResult(interp resolvedPython, outcome pythonExecOutcome, trimmed string, timeoutSeconds int64) string {
336 var sb strings.Builder
337 fmt.Fprintf(&sb, "python %s (%s)\n", interp.Version.String(), interp.Path)
338 if trimmed != "" {
339 sb.WriteString("\n")
340 sb.WriteString(trimmed)
341 if !strings.HasSuffix(trimmed, "\n") {
342 sb.WriteString("\n")
343 }
344 }
345 if outcome.LinesTrunc {
346 fmt.Fprintf(&sb, "... [truncated at %d lines]\n", maxPythonLines)
347 }
348 if outcome.BytesTrunc {
349 fmt.Fprintf(&sb, "... [truncated at %d bytes]\n", maxPythonBytes)
350 }
351 sb.WriteString("\n")
352 switch {
353 case outcome.TimedOut:
354 fmt.Fprintf(&sb, "[timed out after %d second", timeoutSeconds)
355 if timeoutSeconds != 1 {
356 sb.WriteByte('s')
357 }
358 sb.WriteByte(']')
359 case outcome.Cancelled:
360 sb.WriteString("[cancelled]")
361 case outcome.ExitCode == 0:
362 sb.WriteString("[exit 0]")
363 default:
364 fmt.Fprintf(&sb, "[exit %d]", outcome.ExitCode)
365 }
366 if outcome.CleanupErr != nil {
367 fmt.Fprintf(&sb, " [cleanup: %s]", outcome.CleanupErr.Error())
368 }
369 return sb.String()
370}