fix(tools): cap bash stdout/stderr capture to prevent OOM

internal/tools/bash.go

@@ -18,6 +18,40 @@ import (
 const defaultBashTimeoutMS = 120000
 const maxBashTimeoutMS = 600000
 
+const maxBashOutputBytes = 16 * 1024 * 1024 // 16 MiB - prevent OOM on large command output
+
+// cappedWriter wraps a bytes.Buffer and stops accepting bytes after max,
+// setting a flag instead. Used to cap stdout/stderr capture during
+// command.Run() instead of only truncating afterward.
+//
+// The buffer is a named field rather than embedded on purpose: embedding
+// *bytes.Buffer would promote its ReadFrom method, and exec.Cmd's Run()
+// drives non-*os.File Stdout/Stderr writers through io.Copy, which prefers
+// io.ReaderFrom over Write when the destination implements it. That would
+// let the child process's output flow straight into the buffer via the
+// promoted ReadFrom, bypassing the cap entirely and defeating the whole
+// point of this type.
+type cappedWriter struct {
+	buf       *bytes.Buffer
+	max       int
+	truncated *bool
+}
+
+func (c *cappedWriter) Write(p []byte) (int, error) {
+	if c.buf.Len()+len(p) > c.max {
+		*c.truncated = true
+		remain := c.max - c.buf.Len()
+		if remain > 0 {
+			c.buf.Write(p[:remain])
+		}
+		return len(p), nil
+	}
+	return c.buf.Write(p)
+}
+
+func (c *cappedWriter) String() string { return c.buf.String() }
+func (c *cappedWriter) Len() int       { return c.buf.Len() }
+
 type bashTool struct {
 	baseTool
 	workspaceRoot string
@@ -124,10 +158,12 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
 	defer plan.Cleanup()
 	addSandboxMeta(meta, plan)
 
-	var stdout bytes.Buffer
-	var stderr bytes.Buffer
-	command.Stdout = &stdout
-	command.Stderr = &stderr
+	stdoutTruncated := false
+	stderrTruncated := false
+	stdout := &cappedWriter{buf: &bytes.Buffer{}, max: maxBashOutputBytes, truncated: &stdoutTruncated}
+	stderr := &cappedWriter{buf: &bytes.Buffer{}, max: maxBashOutputBytes, truncated: &stderrTruncated}
+	command.Stdout = stdout
+	command.Stderr = stderr
 
 	// Kill the shell as a process group on timeout and bound the post-kill I/O
 	// wait, so a backgrounded child cannot outlive the command or hang Run().
@@ -159,7 +195,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
 		markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText)
 		return Result{
 			Status: StatusError,
-			Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
+			Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta, stdoutTruncated, stderrTruncated),
 			Meta:   meta,
 		}
 	}
@@ -168,13 +204,13 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
 	if meta[SandboxLikelyDeniedMeta] == "true" {
 		return Result{
 			Status: StatusError,
-			Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta),
+			Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta, stdoutTruncated, stderrTruncated),
 			Meta:   meta,
 		}
 	}
 	return Result{
 		Status: StatusOK,
-		Output: formatBashOutput(stdout.String(), stderrText, exitCode),
+		Output: formatBashOutputWithTruncation(stdout.String(), stderrText, exitCode, stdoutTruncated, stderrTruncated),
 		Meta:   meta,
 	}
 }
@@ -319,6 +355,10 @@ func commandExitCode(err error) int {
 }
 
 func formatBashOutput(stdout string, stderr string, exitCode int) string {
+	return formatBashOutputWithTruncation(stdout, stderr, exitCode, false, false)
+}
+
+func formatBashOutputWithTruncation(stdout string, stderr string, exitCode int, stdoutTruncated, stderrTruncated bool) string {
 	parts := []string{}
 	stdout = strings.TrimRight(stdout, "\r\n")
 	stderr = strings.TrimRight(stderr, "\r\n")
@@ -339,14 +379,20 @@ func formatBashOutput(stdout string, stderr string, exitCode int) string {
 	if n := len(outFindings) + len(errFindings); n > 0 {
 		parts = append(parts, fmt.Sprintf("[zero] redacted %d likely secret(s) from this output before showing it.", n))
 	}
+	if stdoutTruncated {
+		parts = append(parts, fmt.Sprintf("[zero] output truncated at %d MiB", maxBashOutputBytes/(1024*1024)))
+	}
+	if stderrTruncated {
+		parts = append(parts, fmt.Sprintf("[zero] stderr truncated at %d MiB", maxBashOutputBytes/(1024*1024)))
+	}
 	if len(parts) == 0 {
 		return "Command completed with no output."
 	}
 	return strings.Join(parts, "\n")
 }
 
-func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string) string {
-	output := formatBashOutput(stdout, stderr, exitCode)
+func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string, stdoutTruncated, stderrTruncated bool) string {
+	output := formatBashOutputWithTruncation(stdout, stderr, exitCode, stdoutTruncated, stderrTruncated)
 	if issue := detectShellOutputIssue(command, stdout+"\n"+stderr, runtime.GOOS); issue != nil {
 		meta["shell_issue"] = issue.Kind
 		output = appendShellIssueHint(output, *issue)

internal/tools/bash_tool_test.go

@@ -1,8 +1,10 @@
 package tools
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"os"
@@ -47,6 +49,13 @@ func runBashToolHelper(command string) {
 	case "long-sleep":
 		time.Sleep(5 * time.Second)
 		fmt.Println("long sleep finished")
+	case "large-output":
+		// Writes well beyond maxBashOutputBytes so tests can verify the
+		// capped writer stops growing memory instead of buffering it all.
+		chunk := bytes.Repeat([]byte("x"), 1024*1024)
+		for i := 0; i < 20; i++ {
+			os.Stdout.Write(chunk)
+		}
 	case "http-server":
 		listener, err := net.Listen("tcp", "127.0.0.1:0")
 		if err != nil {
@@ -298,6 +307,95 @@ func TestBashToolReturnsNonzeroExitAsError(t *testing.T) {
 	}
 }
 
+func TestBashToolCapsLargeStdoutAndReportsTruncation(t *testing.T) {
+	result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{
+		"command": helperCommand("large-output"),
+	})
+
+	if result.Status != StatusOK {
+		t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output)
+	}
+	wantNotice := fmt.Sprintf("output truncated at %d MiB", maxBashOutputBytes/(1024*1024))
+	if !strings.Contains(result.Output, wantNotice) {
+		t.Fatalf("expected truncation notice %q, got output of length %d", wantNotice, len(result.Output))
+	}
+	if len(result.Output) > maxBashOutputBytes+4096 {
+		t.Fatalf("expected output to stay bounded near the cap, got length %d", len(result.Output))
+	}
+}
+
+func TestCappedWriterStopsGrowingAfterMax(t *testing.T) {
+	truncated := false
+	writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated}
+
+	n, err := writer.Write([]byte("0123456789"))
+	if err != nil || n != 10 {
+		t.Fatalf("first write = (%d, %v), want (10, nil)", n, err)
+	}
+	if truncated {
+		t.Fatalf("truncated flag set early, buffer exactly at max")
+	}
+	if writer.Len() != 10 {
+		t.Fatalf("buffer length = %d, want 10", writer.Len())
+	}
+
+	n, err = writer.Write([]byte("overflow"))
+	if err != nil || n != len("overflow") {
+		t.Fatalf("second write = (%d, %v), want (%d, nil)", n, err, len("overflow"))
+	}
+	if !truncated {
+		t.Fatalf("expected truncated flag to be set once max exceeded")
+	}
+	if writer.Len() != 10 {
+		t.Fatalf("buffer should not grow past max, got length %d", writer.Len())
+	}
+}
+
+func TestCappedWriterCapturesPartialWriteThatCrossesMax(t *testing.T) {
+	truncated := false
+	writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated}
+
+	n, err := writer.Write([]byte("012345"))
+	if err != nil || n != 6 {
+		t.Fatalf("first write = (%d, %v), want (6, nil)", n, err)
+	}
+
+	n, err = writer.Write([]byte("6789ABCDEF"))
+	if err != nil || n != 10 {
+		t.Fatalf("second write = (%d, %v), want (10, nil)", n, err)
+	}
+	if !truncated {
+		t.Fatalf("expected truncated flag to be set")
+	}
+	if got, want := writer.String(), "0123456789"; got != want {
+		t.Fatalf("buffer contents = %q, want %q", got, want)
+	}
+}
+
+func TestCappedWriterEnforcesCapThroughIOCopy(t *testing.T) {
+	// exec.Cmd drives a non-*os.File Stdout/Stderr through io.Copy, which
+	// prefers io.ReaderFrom over Write when the destination implements it.
+	// This reproduces that path directly (regression test for the bug where
+	// embedding *bytes.Buffer promoted ReadFrom and bypassed the cap).
+	truncated := false
+	writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated}
+
+	if _, ok := any(writer).(io.ReaderFrom); ok {
+		t.Fatalf("cappedWriter must not implement io.ReaderFrom, or io.Copy bypasses Write and the cap")
+	}
+
+	source := bytes.NewReader([]byte("0123456789overflow"))
+	if _, err := io.Copy(writer, source); err != nil {
+		t.Fatalf("io.Copy returned error: %v", err)
+	}
+	if !truncated {
+		t.Fatalf("expected truncated flag to be set")
+	}
+	if writer.Len() != 10 {
+		t.Fatalf("expected io.Copy through cappedWriter to stay capped at 10 bytes, got %d", writer.Len())
+	}
+}
+
 func TestBashToolTimesOut(t *testing.T) {
 	result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{
 		"command":    helperCommand("sleep"),