fix(output): apply semantic budgets to captured output

Author: KRATOS

internal/tools/bash.go

@@ -71,14 +71,18 @@ func NewScopedBashTool(workspaceRoot string, scope PathScope) Tool {
 }
 
 func (tool bashTool) Run(ctx context.Context, args map[string]any) Result {
-	return tool.run(ctx, args, nil)
+	return tool.run(ctx, args, nil, true)
 }
 
 func (tool bashTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
-	return tool.run(ctx, args, engine)
+	return tool.run(ctx, args, engine, true)
 }
 
-func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
+func (tool bashTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
+	return tool.run(ctx, args, options.Sandbox, false)
+}
+
+func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine, directBudget bool) Result {
 	commandText, err := aliasedStringArg(args, []string{"command", "cmd", "script", "shell"}, "", true, false)
 	if err != nil {
 		return errorResult("Error: Invalid arguments for bash: " + err.Error())
@@ -180,7 +184,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
 			}
 		}
 		markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
-		outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
+		outText, errText, truncated := prepareBashOutput(stdoutText, stdout.total, stderrText, stderrTotal, meta, directBudget)
 		return Result{
 			Status:    StatusError,
 			Output:    formatBashOutputWithShellHint(outText, errText, exitCode, meta),
@@ -190,7 +194,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
 	}
 
 	markLikelySandboxDenial(meta, plan, exitCode, stdoutText, stderrText)
-	outText, errText, truncated := budgetBashCapture(stdoutText, stdout.total, stderrText, stderrTotal, meta)
+	outText, errText, truncated := prepareBashOutput(stdoutText, stdout.total, stderrText, stderrTotal, meta, directBudget)
 	if meta[SandboxLikelyDeniedMeta] == "true" {
 		return Result{
 			Status:    StatusError,
@@ -423,6 +427,28 @@ func budgetBashOutput(stdout string, stderr string, meta map[string]string) (str
 	return budgetBashCapture(stdout, len(stdout), stderr, len(stderr), meta)
 }
 
+// prepareBashOutput keeps direct callers on the established positional budget.
+// Registry calls retain the existing bounded capture but leave final semantic
+// reduction and spill creation to the post-redaction registry boundary.
+func prepareBashOutput(out string, outTotal int, errStr string, errTotal int, meta map[string]string, directBudget bool) (string, string, bool) {
+	if directBudget {
+		return budgetBashCapture(out, outTotal, errStr, errTotal, meta)
+	}
+	outText := sectionWithCaptureGap(out, outTotal)
+	errText := sectionWithCaptureGap(errStr, errTotal)
+	truncated := outTotal > len(out) || errTotal > len(errStr)
+	if meta != nil {
+		meta["raw_bytes"] = strconv.Itoa(outTotal + errTotal)
+		meta["emitted_bytes"] = strconv.Itoa(len(outText) + len(errText))
+		meta["estimated_tokens"] = strconv.Itoa(estimatedTokensFromBytes(len(outText) + len(errText)))
+		if truncated {
+			meta["truncated"] = "true"
+			meta["truncation_reason"] = "capture_budget"
+		}
+	}
+	return outText, errText, truncated
+}
+
 // budgetBashCapture is budgetBashOutput for the streaming-capture path: outTotal
 // and errTotal are the true byte counts (from boundedBuffer.total), which may
 // exceed the retained strings when the middle was dropped during capture. Meta's

internal/tools/bash_budget_test.go

@@ -37,6 +37,18 @@ func TestBudgetBashOutputSmallPassesThrough(t *testing.T) {
 	}
 }
 
+func TestPrepareBashOutputDefersRegistrySemanticBudget(t *testing.T) {
+	output := "START\n" + strings.Repeat("progress\n", 8_000) + "ERROR: retained\nEND"
+	registryOutput, _, registryTruncated := prepareBashOutput(output, len(output), "", 0, map[string]string{}, false)
+	if registryTruncated || registryOutput != output {
+		t.Fatalf("registry preparation must preserve capture-bounded output for semantic reduction: truncated=%t bytes=%d", registryTruncated, len(registryOutput))
+	}
+	directOutput, _, directTruncated := prepareBashOutput(output, len(output), "", 0, map[string]string{}, true)
+	if !directTruncated || len(directOutput) > bashOutputBudgetBytes+512 {
+		t.Fatalf("direct path must preserve legacy positional budget: truncated=%t bytes=%d", directTruncated, len(directOutput))
+	}
+}
+
 // Oversized stdout is truncated head+tail: both the first and last lines survive,
 // the middle is dropped behind a marker, meta is flagged, and the captured text
 // is spilled to a re-readable file.

internal/tools/exec_command.go

@@ -483,11 +483,15 @@ func NewScopedExecCommandTool(workspaceRoot string, scope PathScope, manager *ex
 }
 
 func (tool execCommandTool) Run(ctx context.Context, args map[string]any) Result {
-	return tool.run(ctx, args, nil)
+	return tool.run(ctx, args, nil, true)
 }
 
 func (tool execCommandTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
-	return tool.run(ctx, args, engine)
+	return tool.run(ctx, args, engine, true)
+}
+
+func (tool execCommandTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
+	return tool.run(ctx, args, options.Sandbox, false)
 }
 
 func (tool execCommandTool) ExecSessions() []ExecSessionSnapshot {
@@ -502,7 +506,7 @@ func (tool execCommandTool) StopAllExecSessions() []int {
 	return tool.manager.stopAll()
 }
 
-func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
+func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine, directBudget bool) Result {
 	commandText, err := aliasedStringArg(args, []string{"cmd", "command", "script", "shell"}, "", true, false)
 	if err != nil {
 		return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
@@ -558,7 +562,7 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
 	if exited {
 		tool.manager.remove(session.id)
 	}
-	return execToolResult(execToolResultInput{
+	return execToolResultWithBudget(execToolResultInput{
 		commandText:           commandText,
 		output:                output,
 		outputBufferTruncated: outputTruncated,
@@ -569,7 +573,7 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
 		tty:                   session.tty,
 		plan:                  session.plan,
 		maxOutputTokens:       maxOutputTokens,
-	})
+	}, directBudget)
 }
 
 func (tool execCommandTool) startSession(commandText string, absoluteCwd string, relativeCwd string, ttyRequested bool, engine *zeroSandbox.Engine, sandboxPermissions SandboxPermissionOverride) (*execSession, error) {
@@ -855,7 +859,15 @@ type execToolResultInput struct {
 }
 
 func execToolResult(input execToolResultInput) Result {
-	output, truncated := truncateExecOutput(input.output, input.maxOutputTokens)
+	return execToolResultWithBudget(input, true)
+}
+
+func execToolResultWithBudget(input execToolResultInput, directBudget bool) Result {
+	output := input.output
+	truncated := false
+	if directBudget {
+		output, truncated = truncateExecOutput(input.output, input.maxOutputTokens)
+	}
 	meta := map[string]string{
 		"cwd": input.relativeCwd,
 		"tty": strconv.FormatBool(input.tty),

internal/tools/output_boundary.go

@@ -34,6 +34,50 @@ func applyRegistryOutputBudget(tool Tool, toolName string, args map[string]any,
 	return result
 }
 
+// applySelfManagedOutputBudget applies semantic retention to tools that retain
+// output with their own capture-aware limits. Their explicit budgets remain
+// authoritative, rather than being replaced by the registry default ceiling.
+func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]any, result Result) Result {
+	budget := selfManagedOutputBudget(toolName, args)
+	if budget.maxEstimatedTokens <= 0 && budget.hardMaxBytes <= 0 {
+		return annotateSelfBudgetedOutput(tool, toolName, args, result)
+	}
+
+	category := resolveOutputCategory(tool, toolName, args)
+	budgeted := budgetSemanticOutput(result.Output, category, budget)
+	if result.Truncated && !budgeted.truncated {
+		budgeted.truncated = true
+		budgeted.reason = result.Meta["truncation_reason"]
+		if budgeted.reason == "" {
+			budgeted.reason = "upstream_tool_budget"
+		}
+	}
+	if budgeted.truncated {
+		budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
+	}
+	result.Output = budgeted.text
+	result.Truncated = result.Truncated || budgeted.truncated
+	result.Meta = addOutputBudgetMetadata(result.Meta, budgeted)
+	return result
+}
+
+func selfManagedOutputBudget(toolName string, args map[string]any) outputBudget {
+	switch toolName {
+	case "bash":
+		// bash previously exposed at most 32 KiB from each stream. Keep that
+		// combined 64 KiB ceiling while selecting its contents semantically.
+		return outputBudget{maxEstimatedTokens: (bashOutputBudgetBytes * 2) / 4, hardMaxBytes: bashOutputBudgetBytes * 2}
+	case ExecCommandToolName:
+		maxOutputTokens := defaultMaxOutputTokens
+		if parsed, err := intArg(args, "max_output_tokens", defaultMaxOutputTokens, 1, maxExecOutputTokenRequest); err == nil {
+			maxOutputTokens = parsed
+		}
+		return outputBudget{maxEstimatedTokens: maxOutputTokens, hardMaxBytes: maxOutputTokens * 4}
+	default:
+		return outputBudget{}
+	}
+}
+
 // RebudgetAfterHook reapplies the registry's redaction and output limits after
 // an afterTool hook appends model-visible feedback to a completed result. The
 // initial registry pass has already run; this second pass is limited to the
@@ -150,6 +194,11 @@ func attachExistingSpill(toolName, output string, budget outputBudget, current b
 	base.text = text
 	base.retainedBytes = len(text)
 	base.estimatedRetainedTokens = estimateOutputTokens(text)
+	base.truncated = true
+	if base.reason == "" {
+		base.reason = current.reason
+	}
+	base.category = current.category
 	base.spillPath = path
 	return base
 }

internal/tools/output_boundary_test.go

@@ -163,6 +163,9 @@ func TestRegistryLargeSingleLineFileRetainsHeadAndTail(t *testing.T) {
 	if len(result.Output) > readOutputBudgetBytes || !utf8.ValidString(result.Output) {
 		t.Fatalf("single-line file output is not safely bounded: bytes=%d valid=%t", len(result.Output), utf8.ValidString(result.Output))
 	}
+	if rawBytes, err := strconv.Atoi(result.Meta["raw_bytes"]); err != nil || rawBytes <= len(result.Output) {
+		t.Fatalf("single-line file capture was not bounded head/tail output: raw=%q emitted=%d", result.Meta["raw_bytes"], len(result.Output))
+	}
 	for _, want := range []string{"File: large.min.js", "HEAD_SINGLE_LINE_MARK", "TAIL_SINGLE_LINE_MARK"} {
 		if !strings.Contains(result.Output, want) {
 			t.Fatalf("single-line file output missing %q", want)
@@ -235,13 +238,72 @@ func TestShellOutputCategoryClassification(t *testing.T) {
 	}
 }
 
-func TestSelfBudgetedToolDeclaresSemanticCategoryWithoutRebudgeting(t *testing.T) {
+func TestSelfManagedOutputBudgetUsesSemanticTestPolicy(t *testing.T) {
+	setTestTempDir(t)
+	input := "stdout:\n" + strings.Repeat("PASS progress\n", 8_000) +
+		"--- FAIL: TestImportant (0.01s)\nthing_test.go:42: expected 7, got 9\nFAIL\nexit_code: 1"
+	tests := []struct {
+		name string
+		tool Tool
+		args map[string]any
+	}{
+		{name: "bash", tool: NewBashTool(t.TempDir()), args: map[string]any{"command": "go test ./..."}},
+		{name: "exec", tool: NewExecCommandTool(t.TempDir(), newExecSessionManager()), args: map[string]any{"cmd": "go test ./...", "max_output_tokens": 16_000}},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := applySelfManagedOutputBudget(test.tool, test.tool.Name(), test.args, Result{Status: StatusError, Output: input, Meta: map[string]string{}})
+			if !got.Truncated || got.Meta[outputBudgetCategoryMeta] != string(outputCategoryTest) {
+				t.Fatalf("semantic budget was not applied: %#v", got)
+			}
+			for _, want := range []string{"TestImportant", "expected 7", "FAIL", "exit_code: 1"} {
+				if !strings.Contains(got.Output, want) {
+					t.Fatalf("semantic test output missing %q:\n%s", want, got.Output)
+				}
+			}
+		})
+	}
+}
+
+type semanticBashFakeTool struct{ ceilingFakeTool }
+
+func (semanticBashFakeTool) managesOutputBudget() {}
+
+func (semanticBashFakeTool) outputCategory(args map[string]any) outputCategory {
+	command, _ := args["command"].(string)
+	return shellOutputCategory(command)
+}
+
+func TestRegistryAppliesSemanticBudgetToSelfManagedShellOutput(t *testing.T) {
+	setTestTempDir(t)
+	input := "startup\n" + strings.Repeat("PASS progress\n", 8_000) +
+		"--- FAIL: TestRegistryImportant (0.01s)\nexpected 7, got 9\nFAIL\nexit_code: 1"
+	registry := NewRegistry()
+	registry.Register(semanticBashFakeTool{newCeilingFakeTool("bash", input)})
+	result := registry.Run(context.Background(), "bash", map[string]any{"command": "go test ./..."})
+	if !result.Truncated || result.Meta[outputBudgetReasonMeta] != "semantic_test_budget" {
+		t.Fatalf("registry did not apply semantic test budget: %#v", result.Meta)
+	}
+	for _, want := range []string{"TestRegistryImportant", "expected 7", "FAIL", "exit_code: 1"} {
+		if !strings.Contains(result.Output, want) {
+			t.Fatalf("registry semantic test output missing %q:\n%s", want, result.Output)
+		}
+	}
+}
+
+func TestRegistryAppliesSemanticBudgetToSelfManagedProcessOutput(t *testing.T) {
+	setTestTempDir(t)
+	input := "starting build\n" + strings.Repeat("progress\n", 8_000) +
+		"WARNING: cache is cold\nERROR: compilation failed\nexit_code: 1"
 	registry := NewRegistry()
-	registry.Register(NewBashTool(t.TempDir()))
-	tool, _ := registry.Get("bash")
-	result := Result{Status: StatusOK, Output: "already bounded", Meta: map[string]string{}}
-	got := annotateSelfBudgetedOutput(tool, "bash", map[string]any{"command": "go test ./..."}, result)
-	if got.Output != result.Output || got.Meta[outputBudgetCategoryMeta] != string(outputCategoryTest) {
-		t.Fatalf("self-budgeted category annotation changed output or lost category: %#v", got)
+	registry.Register(semanticBashFakeTool{newCeilingFakeTool("bash", input)})
+	result := registry.Run(context.Background(), "bash", map[string]any{"command": "make build"})
+	if !result.Truncated || result.Meta[outputBudgetReasonMeta] != "semantic_process_budget" {
+		t.Fatalf("registry did not apply semantic process budget: %#v", result.Meta)
+	}
+	for _, want := range []string{"starting build", "WARNING", "ERROR", "exit_code: 1"} {
+		if !strings.Contains(result.Output, want) {
+			t.Fatalf("registry semantic process output missing %q:\n%s", want, result.Output)
+		}
 	}
 }

internal/tools/output_budget.go

@@ -218,6 +218,7 @@ func estimatedTokensFromBytes(bytes int) int {
 
 type outputBudgetBuilder struct {
 	builder  strings.Builder
+	capture  *boundedBuffer
 	rawBytes int
 	maxBytes int
 	hint     string
@@ -227,8 +228,24 @@ func newOutputBudgetBuilder(maxBytes int, hint string) *outputBudgetBuilder {
 	return &outputBudgetBuilder{maxBytes: maxBytes, hint: hint}
 }
 
+// newHeadTailOutputBudgetBuilder bounds capture while retaining both ends of
+// an input for a later semantic policy. It is used for registry-driven file
+// reads; direct Tool.Run calls keep the legacy prefix-only builder above.
+func newHeadTailOutputBudgetBuilder(maxBytes int, hint string) *outputBudgetBuilder {
+	headBytes := maxBytes / 2
+	return &outputBudgetBuilder{
+		capture:  newBoundedBuffer(headBytes, maxBytes-headBytes),
+		maxBytes: maxBytes,
+		hint:     hint,
+	}
+}
+
 func (builder *outputBudgetBuilder) WriteString(value string) {
 	builder.rawBytes += len(value)
+	if builder.capture != nil {
+		_, _ = builder.capture.Write([]byte(value))
+		return
+	}
 	if builder.maxBytes <= 0 {
 		builder.builder.WriteString(value)
 		return
@@ -262,6 +279,9 @@ func applyLegacyByteBudgetToResult(result Result, maxBytes int, hint string) Res
 
 func (builder *outputBudgetBuilder) Result() outputBudgetResult {
 	output := builder.builder.String()
+	if builder.capture != nil {
+		output = builder.capture.retained()
+	}
 	result := outputBudgetResult{
 		Output:       output,
 		RawBytes:     builder.rawBytes,
@@ -276,7 +296,12 @@ func (builder *outputBudgetBuilder) Result() outputBudgetResult {
 	if budget < 0 {
 		budget = 0
 	}
-	result.Output = utf8Prefix(output, budget) + marker
+	if builder.capture != nil {
+		headBytes := budget / 2
+		result.Output = utf8Prefix(output, headBytes) + marker + utf8Suffix(output, budget-headBytes)
+	} else {
+		result.Output = utf8Prefix(output, budget) + marker
+	}
 	result.Truncated = true
 	result.EmittedBytes = len(result.Output)
 	return result

internal/tools/output_ceiling.go

@@ -25,16 +25,17 @@ const defaultOutputCeilingTokens = 16_000
 // disables the ceiling entirely; unset or unparsable keeps the default.
 const outputCeilingEnv = "ZERO_TOOL_OUTPUT_CEILING_TOKENS"
 
-// selfBudgeting marks a tool that enforces its own deliberate output budget —
-// possibly model-raisable (exec_command) — which the registry ceiling must not
-// second-guess. The method is unexported on purpose: only tools in this
-// package can opt out, so an MCP-served tool can never exempt itself.
+// selfBudgeting marks a tool with a deliberate capture-aware output budget —
+// possibly model-raisable (exec_command). The registry applies semantic
+// retention within that explicit budget rather than replacing it with the
+// default ceiling. The method is unexported so an MCP-served tool cannot opt
+// into this path.
 type selfBudgeting interface{ managesOutputBudget() }
 
-// The exemption list, kept in one place. Shell/process tools retain their
-// established capture-aware budgets: bash (per-stream budget + spill) and
-// exec_command (model-raisable budget + spill). File/search tools now use the
-// shared post-redaction semantic boundary.
+// The list is kept in one place. Shell/process tools retain their established
+// capture-aware budgets: bash (per-stream bounded capture) and exec_command
+// (model-raisable bounded session output). File/search tools use the standard
+// shared post-redaction boundary.
 func (bashTool) managesOutputBudget()        {}
 func (execCommandTool) managesOutputBudget() {}