Compare commits

...
30 changed files with 2225 additions and 121 deletions

No files matched your search

+16 -3
View File
@@ -272,7 +272,8 @@ sequenceDiagram
Agent->>Tool: execute(ctx, args)
Tool-->>Agent: output, error, metadata
Agent->>Hooks: afterTool hooks
Agent->>Agent: redact, truncate, classify success/failure
Agent->>Agent: redact secrets, apply semantic output budget, enforce byte ceiling
Agent->>Agent: classify success/failure
Agent-->>Surface: OnToolResult callback
Agent->>Transcript: append tool result message
end
@@ -604,10 +605,22 @@ flowchart TD
Prompt -- Yes --> Decision[Permission callback]
Prompt -- No --> Run[Registry.RunWithOptions]
Decision --> Run
Run --> Redact[Redact secrets + enforce output ceiling]
Redact --> Message[Return tool result to model]
Run --> Redact[Redact secrets]
Redact --> Budget[Token-aware semantic output budget]
Budget --> Ceiling[Existing hard byte ceiling]
Ceiling --> Message[Return one tool result to model]
```
Oversized results use deterministic, provider-neutral estimated-token budgets.
Policies retain useful structure for files, search matches, tests, process logs,
diffs, and worker conclusions; tools without a declared category use a UTF-8-safe
head/tail fallback. The estimate is intentionally conservative for non-ASCII
text and is not exact provider tokenization. Existing byte ceilings remain the
authoritative safety limit. When the existing spill mechanism can persist the
complete redacted text received by the budgeting layer, the result includes its
safe spill reference. This does not imply capture of subprocess bytes already
discarded by a tool's established internal buffer.
Core tool groups include:
- **Read-only tools**: file reads, directory listing, glob, grep, LSP navigation,
+6 -1
View File
@@ -141,9 +141,11 @@ func TestRunAskUserCancellationAbortsRun(t *testing.T) {
registry := registryWithAskUser()
args := `{"questions":[{"question":"Which framework?"}]}`
provider := providerCallingAskUserThenAnswer(args, "done")
var toolResults []ToolResult
result, err := Run(context.Background(), "clarify", provider, Options{
Registry: registry,
Registry: registry,
OnToolResult: func(result ToolResult) { toolResults = append(toolResults, result) },
OnAskUser: func(_ context.Context, _ AskUserRequest) (AskUserResponse, error) {
return AskUserResponse{}, context.Canceled
},
@@ -158,6 +160,9 @@ func TestRunAskUserCancellationAbortsRun(t *testing.T) {
if len(provider.requests) != 1 {
t.Fatalf("expected the run to stop after the canceled ask_user (1 turn), got %d", len(provider.requests))
}
if len(toolResults) != 1 || toolResults[0].ToolCallID == "" {
t.Fatalf("cancellation must emit exactly one result for the call, got %#v", toolResults)
}
// The recorded tool result must reflect cancellation, not a synthetic answer.
var toolMsg string
for _, m := range result.Messages {
+28
View File
@@ -7,6 +7,7 @@ import (
"fmt"
"io"
"sort"
"strconv"
"strings"
"sync"
@@ -656,6 +657,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
toolSpan.End()
}
options.Trace.Counter(trace.CounterToolCalls, 1)
recordOutputBudgetTrace(options.Trace, toolResult)
if options.OnToolResult != nil {
options.OnToolResult(toolResult)
}
@@ -841,6 +843,28 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
return result, nil
}
func recordOutputBudgetTrace(recorder *trace.Recorder, result ToolResult) {
if recorder == nil || result.Meta["output_budget_category"] == "" {
return
}
parseInt := func(key string) int {
value, _ := strconv.Atoi(result.Meta[key])
return value
}
spillCreated, _ := strconv.ParseBool(result.Meta["output_budget_spill_created"])
recorder.EmitOutputBudget(trace.OutputBudgetEvent{
Tool: result.Name,
Category: result.Meta["output_budget_category"],
OriginalBytes: parseInt("output_budget_original_bytes"),
RetainedBytes: parseInt("output_budget_retained_bytes"),
EstimatedOriginalTokens: parseInt("output_budget_estimated_original_tokens"),
EstimatedRetainedTokens: parseInt("output_budget_estimated_retained_tokens"),
Truncated: result.Truncated,
Reason: result.Meta["output_budget_reason"],
SpillCreated: spillCreated,
})
}
func finalAnswerAfterMaxTurns(ctx context.Context, provider Provider, messages []zeroruntime.Message, toolDefs []zeroruntime.ToolDefinition, options Options) (string, []zeroruntime.Message, string) {
finalMessages := copyMessages(messages)
finalMessages = append(finalMessages, zeroruntime.Message{
@@ -1295,6 +1319,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
if didRedact {
result.Redacted = true
}
result = registry.RebudgetAfterHook(call.Name, args, result)
}
}
// Secret scrubbing happens at the registry boundary (the single point both
@@ -1305,6 +1330,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
Name: call.Name,
Status: result.Status,
Output: result.Output,
Truncated: result.Truncated,
Meta: result.Meta,
Redacted: result.Redacted,
ChangedFiles: result.ChangedFiles,
@@ -1615,6 +1641,7 @@ func toolResultFromPrePermissionReject(call ToolCall, result tools.Result) ToolR
Name: call.Name,
Status: result.Status,
Output: output,
Truncated: result.Truncated,
Meta: meta,
Redacted: result.Redacted || outputRedacted || summaryRedacted || metaRedacted,
ChangedFiles: result.ChangedFiles,
@@ -1884,6 +1911,7 @@ func askUserFallbackResult(ctx context.Context, registry *tools.Registry, call T
Name: call.Name,
Status: result.Status,
Output: result.Output,
Truncated: result.Truncated,
Meta: result.Meta,
Redacted: result.Redacted,
ChangedFiles: result.ChangedFiles,
@@ -0,0 +1,167 @@
package agent
import (
"context"
"os"
"strconv"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/hooks"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
type propagationOutputTool struct {
output string
}
func TestOutputBudgetHookHelperProcess(t *testing.T) {
for index, arg := range os.Args {
if arg != "--zero-output-budget-hook" || index+1 >= len(os.Args) {
continue
}
if _, err := os.Stdout.WriteString(os.Args[index+1]); err != nil {
os.Exit(2)
}
os.Exit(0)
}
}
func largeOutputBudgetHookDispatcher() *hooks.Dispatcher {
feedback := strings.Repeat("hook feedback ", 200)
return hooks.NewDispatcher(hooks.DispatcherOptions{Config: hooks.Config{
Enabled: true,
Hooks: []hooks.Definition{{
ID: "large-feedback",
Event: hooks.EventAfterTool,
Matcher: "propagation_output",
Command: os.Args[0],
Args: []string{
"-test.run=TestOutputBudgetHookHelperProcess",
"--",
"--zero-output-budget-hook",
feedback,
},
Enabled: true,
}},
}})
}
func (tool propagationOutputTool) Name() string { return "propagation_output" }
func (tool propagationOutputTool) Description() string { return "returns output for propagation tests" }
func (tool propagationOutputTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", AdditionalProperties: false}
}
func (tool propagationOutputTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "test read"}
}
func (tool propagationOutputTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: tool.output}
}
func TestExecuteToolCallPropagatesOutputTruncation(t *testing.T) {
t.Setenv("TMPDIR", t.TempDir())
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
registry := tools.NewRegistry()
registry.Register(propagationOutputTool{output: strings.Repeat("large output\n", 1000)})
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
ID: "call-budget",
Name: "propagation_output",
Arguments: `{}`,
}, PermissionModeAuto, Options{Cwd: t.TempDir()})
if abortErr != nil {
t.Fatalf("executeToolCall abort error: %v", abortErr)
}
if !result.Truncated {
t.Fatalf("agent ToolResult lost tools.Result.Truncated: %#v", result)
}
if result.Meta["spill_path"] == "" {
t.Fatalf("agent ToolResult lost spill metadata: %#v", result.Meta)
}
}
func TestRecordOutputBudgetTraceUsesOnlyCompactMetadata(t *testing.T) {
recorder := trace.NewRecorder("session", "run", "")
recorder.Start()
recordOutputBudgetTrace(recorder, ToolResult{
Name: "grep",
Truncated: true,
Output: "SECRET OUTPUT MUST NOT ENTER TRACE",
Meta: map[string]string{
"output_budget_category": "search",
"output_budget_original_bytes": "1000",
"output_budget_retained_bytes": "100",
"output_budget_estimated_original_tokens": "250",
"output_budget_estimated_retained_tokens": "25",
"output_budget_reason": "semantic_search_budget",
"output_budget_spill_created": "true",
"spill_path": "/secret/path/not-for-trace",
},
})
events := recorder.Finish().OutputBudgets
if len(events) != 1 {
t.Fatalf("events = %#v", events)
}
event := events[0]
if event.Tool != "grep" || event.Category != "search" || event.OriginalBytes != 1000 || event.RetainedBytes != 100 || !event.SpillCreated {
t.Fatalf("unexpected trace event: %#v", event)
}
}
func TestExecuteToolCallRebudgetsOversizedAfterToolFeedback(t *testing.T) {
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
registry := tools.NewRegistry()
registry.Register(propagationOutputTool{output: "tool output"})
dispatcher := largeOutputBudgetHookDispatcher()
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
ID: "call-hook-budget",
Name: "propagation_output",
Arguments: `{}`,
}, PermissionModeAuto, Options{Hooks: dispatcher})
if abortErr != nil {
t.Fatalf("executeToolCall abort error: %v", abortErr)
}
if !result.Truncated || len(result.Output) > 80*4 {
t.Fatalf("afterTool feedback bypassed output budget: truncated=%t bytes=%d meta=%#v", result.Truncated, len(result.Output), result.Meta)
}
if result.Meta["output_budget_category"] == "" || result.Meta["output_budget_retained_bytes"] != strconv.Itoa(len(result.Output)) {
t.Fatalf("post-hook budget metadata does not describe final output: %#v", result.Meta)
}
}
func TestRunTraceReflectsPostHookBudget(t *testing.T) {
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
registry := tools.NewRegistry()
registry.Register(propagationOutputTool{output: "tool output"})
dispatcher := largeOutputBudgetHookDispatcher()
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
{
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-hook-trace", ToolName: "propagation_output"},
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-hook-trace", ArgumentsFragment: `{}`},
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-hook-trace"},
{Type: zeroruntime.StreamEventDone},
},
{{Type: zeroruntime.StreamEventText, Content: "done"}, {Type: zeroruntime.StreamEventDone}},
}}
recorder := trace.NewRecorder("session", "run", "")
var toolResults []ToolResult
if _, err := Run(context.Background(), "budget hook", provider, Options{
Registry: registry,
Hooks: dispatcher,
Trace: recorder,
OnToolResult: func(result ToolResult) { toolResults = append(toolResults, result) },
}); err != nil {
t.Fatalf("Run: %v", err)
}
if len(toolResults) != 1 || !toolResults[0].Truncated {
t.Fatalf("tool result = %#v, want one truncated post-hook result", toolResults)
}
events := recorder.Finish().OutputBudgets
if len(events) != 1 || !events[0].Truncated || events[0].RetainedBytes != len(toolResults[0].Output) {
t.Fatalf("trace does not describe final post-hook output: events=%#v result=%#v", events, toolResults[0])
}
}
+7 -4
View File
@@ -61,10 +61,13 @@ const (
)
type ToolResult struct {
ToolCallID string
Name string
Status tools.Status
Output string
ToolCallID string
Name string
Status tools.Status
Output string
// Truncated reports that the tool's model-visible output omitted content.
// The full result may be recoverable through Meta["spill_path"].
Truncated bool
Meta map[string]string
Redacted bool
ChangedFiles []string
+3
View File
@@ -636,6 +636,9 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
if len(result.Meta) > 0 {
payload["meta"] = result.Meta
}
if result.Truncated {
payload["truncated"] = true
}
if result.Redacted {
payload["redacted"] = true
}
+5 -1
View File
@@ -153,6 +153,9 @@ func (writer *execEventWriter) toolResult(result agent.ToolResult) {
if len(result.Meta) > 0 {
payload["meta"] = result.Meta
}
if result.Truncated {
payload["truncated"] = true
}
if result.Redacted {
payload["redacted"] = true
}
@@ -166,7 +169,8 @@ func (writer *execEventWriter) toolResult(result agent.ToolResult) {
return
}
if writer.format == execOutputStreamJSON {
output, truncated := truncateForStreamJSONOutput(result.Output)
output, surfaceTruncated := truncateForStreamJSONOutput(result.Output)
truncated := result.Truncated || surfaceTruncated
event := streamjson.Event{
Type: streamjson.EventToolResult,
RunID: writer.runID,
+36
View File
@@ -1,8 +1,12 @@
package cli
import (
"bytes"
"encoding/json"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/tools"
)
@@ -16,3 +20,35 @@ func TestStreamJSONSideEffectReportsNoneForControlTool(t *testing.T) {
t.Fatalf("streamJSONSideEffect(escalate_model) = %q, want none", got)
}
}
func TestExecWriterPropagatesToolResultTruncation(t *testing.T) {
for _, format := range []execOutputFormat{execOutputJSON, execOutputStreamJSON} {
t.Run(string(format), func(t *testing.T) {
var stdout, stderr bytes.Buffer
writer := execEventWriter{
stdout: &stdout,
stderr: &stderr,
format: format,
runID: "run_budget",
streamedText: &strings.Builder{},
}
writer.toolResult(agent.ToolResult{
ToolCallID: "call_budget",
Name: "read_file",
Status: tools.StatusOK,
Output: "bounded output",
Truncated: true,
})
if writer.err != nil {
t.Fatalf("toolResult: %v", writer.err)
}
var payload map[string]any
if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &payload); err != nil {
t.Fatalf("decode output %q: %v", stdout.String(), err)
}
if payload["truncated"] != true {
t.Fatalf("truncated = %#v, want true; payload=%#v", payload["truncated"], payload)
}
})
}
}
+41 -6
View File
@@ -23,6 +23,15 @@ type bashTool struct {
scope PathScope
}
func (bashTool) outputCategory(args map[string]any) outputCategory {
command, _ := bashCommandArg(args)
return shellOutputCategory(command)
}
func bashCommandArg(args map[string]any) (string, error) {
return aliasedStringArg(args, []string{"command", "cmd", "script", "shell"}, "", true, false)
}
func NewBashTool(workspaceRoot string) Tool {
return NewScopedBashTool(workspaceRoot, nil)
}
@@ -63,15 +72,19 @@ 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 {
commandText, err := aliasedStringArg(args, []string{"command", "cmd", "script", "shell"}, "", true, false)
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 := bashCommandArg(args)
if err != nil {
return errorResult("Error: Invalid arguments for bash: " + err.Error())
}
@@ -172,7 +185,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),
@@ -182,7 +195,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,
@@ -415,6 +428,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
+12
View File
@@ -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.
+32 -7
View File
@@ -427,6 +427,15 @@ type execCommandTool struct {
manager *execSessionManager
}
func (execCommandTool) outputCategory(args map[string]any) outputCategory {
command, _ := execCommandArg(args)
return shellOutputCategory(command)
}
func execCommandArg(args map[string]any) (string, error) {
return aliasedStringArg(args, []string{"cmd", "command", "script", "shell"}, "", true, false)
}
func NewExecCommandTool(workspaceRoot string, manager *execSessionManager) Tool {
return NewScopedExecCommandTool(workspaceRoot, nil, manager)
}
@@ -475,11 +484,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 {
@@ -494,8 +507,8 @@ func (tool execCommandTool) StopAllExecSessions() []int {
return tool.manager.stopAll()
}
func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine) Result {
commandText, err := aliasedStringArg(args, []string{"cmd", "command", "script", "shell"}, "", true, false)
func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine *zeroSandbox.Engine, directBudget bool) Result {
commandText, err := execCommandArg(args)
if err != nil {
return errorResult("Error: Invalid arguments for exec_command: " + err.Error())
}
@@ -550,7 +563,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,
@@ -561,7 +574,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) {
@@ -694,6 +707,10 @@ type writeStdinTool struct {
manager *execSessionManager
}
func (writeStdinTool) outputCategory(map[string]any) outputCategory {
return outputCategoryProcess
}
func NewWriteStdinTool(manager *execSessionManager) Tool {
if manager == nil {
manager = defaultExecSessionManager
@@ -843,7 +860,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),
+23 -14
View File
@@ -18,6 +18,8 @@ type globTool struct {
scope PathScope
}
func (globTool) outputCategory(map[string]any) outputCategory { return outputCategorySearch }
func NewGlobTool(workspaceRoot string) Tool {
return NewScopedGlobTool(workspaceRoot, nil)
}
@@ -47,17 +49,25 @@ func NewScopedGlobTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool globTool) Run(ctx context.Context, args map[string]any) Result {
return tool.runWith(ctx, args, readExcluder{})
return tool.runWith(ctx, args, readExcluder{}, true)
}
func (tool globTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
exclude := readExcluder{}
if options.Sandbox != nil {
exclude = sandboxReadExcluder(options.Sandbox)
}
return tool.runWith(ctx, args, exclude, false)
}
// RunWithSandbox runs glob while skipping subtrees the sandbox policy denies
// reads to (DenyRead). With no DenyRead configured the excluder is a no-op and
// behavior is unchanged.
func (tool globTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
return tool.runWith(ctx, args, sandboxReadExcluder(engine), true)
}
func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder, directBudget bool) Result {
pattern, err := aliasedStringArg(args, []string{"pattern", "glob", "match", "query", "expression"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for glob: " + err.Error())
@@ -110,24 +120,23 @@ func (tool globTool) runWith(ctx context.Context, args map[string]any, exclude r
if truncated {
output += fmt.Sprintf("\n\n[truncated: showing first %d of %d matches; increase limit or narrow cwd/pattern]", len(matches), totalMatches)
}
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "increase limit or narrow cwd/pattern")
meta := outputBudgetMeta(budgeted)
meta := map[string]string{}
meta["pattern"] = pattern
if truncated || budgeted.Truncated {
if truncated {
meta["truncated"] = "true"
if budgeted.Truncated {
meta["truncation_reason"] = "byte_budget"
} else {
meta["truncation_reason"] = "limit"
}
meta["truncation_reason"] = "limit"
}
return Result{
result := Result{
Status: StatusOK,
Output: budgeted.Output,
Truncated: truncated || budgeted.Truncated,
Output: output,
Truncated: truncated,
Meta: meta,
}
if directBudget {
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, "increase limit or narrow cwd/pattern")
}
return result
}
func scanGlob(ctx context.Context, root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool, exclude readExcluder) ([]string, error) {
+29 -23
View File
@@ -21,6 +21,8 @@ type grepTool struct {
scope PathScope
}
func (grepTool) outputCategory(map[string]any) outputCategory { return outputCategorySearch }
type grepMatch struct {
file string
line int
@@ -62,7 +64,15 @@ func NewScopedGrepTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool grepTool) Run(ctx context.Context, args map[string]any) Result {
return tool.runWith(ctx, args, readExcluder{})
return tool.runWith(ctx, args, readExcluder{}, true)
}
func (tool grepTool) RunWithOptions(ctx context.Context, args map[string]any, options RunOptions) Result {
exclude := readExcluder{}
if options.Sandbox != nil {
exclude = sandboxReadExcluder(options.Sandbox)
}
return tool.runWith(ctx, args, exclude, false)
}
// RunWithSandbox runs the search while skipping subtrees the sandbox policy
@@ -70,10 +80,10 @@ func (tool grepTool) Run(ctx context.Context, args map[string]any) Result {
// path. With no DenyRead configured the excluder is a no-op and behavior is
// unchanged.
func (tool grepTool) RunWithSandbox(ctx context.Context, args map[string]any, engine *sandbox.Engine) Result {
return tool.runWith(ctx, args, sandboxReadExcluder(engine))
return tool.runWith(ctx, args, sandboxReadExcluder(engine), true)
}
func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder) Result {
func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude readExcluder, directBudget bool) Result {
pattern, err := aliasedStringArg(args, []string{"pattern", "query", "regex", "search", "expression"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for grep: " + err.Error())
@@ -154,7 +164,7 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
}
return errorResult("Error running grep: " + err.Error())
}
return collector.result()
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern to continue")
case "files_with_matches":
collector := &grepFileListCollector{}
if err := scanGrepMatches(ctx, resolvedRoot, target, globMatcher, exclude, absolutePaths, presenceGrepLineMatcher(compiled), collector.collect); err != nil {
@@ -163,7 +173,7 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
}
return errorResult("Error running grep: " + err.Error())
}
return collector.result()
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern to continue")
default:
collector := &grepContentCollector{headLimit: headLimit}
if err := scanGrepMatches(ctx, resolvedRoot, target, globMatcher, exclude, absolutePaths, presenceGrepLineMatcher(compiled), collector.collect); err != nil {
@@ -172,10 +182,17 @@ func (tool grepTool) runWith(ctx context.Context, args map[string]any, exclude r
}
return errorResult("Error running grep: " + err.Error())
}
return collector.result()
return applyDirectSearchBudget(collector.result(), directBudget, "narrow path/glob/pattern or increase head_limit")
}
}
func applyDirectSearchBudget(result Result, directBudget bool, hint string) Result {
if !directBudget {
return result
}
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, hint)
}
// resolveGrepRoot picks the scope root whose EvalSymlinks-resolved path contains
// the already-resolved target, so that confineGrepFile computes correct
// workspace-relative paths even when the target lives in an extra root.
@@ -455,13 +472,7 @@ func (collector *grepFileListCollector) result() Result {
return okResult("No matches found.")
}
sort.Strings(collector.files)
budgeted := applyOutputBudget(strings.Join(collector.files, "\n"), searchOutputBudgetBytes, "narrow path/glob/pattern to continue")
meta := outputBudgetMeta(budgeted)
if budgeted.Truncated {
meta["truncated"] = "true"
meta["truncation_reason"] = "byte_budget"
}
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
return Result{Status: StatusOK, Output: strings.Join(collector.files, "\n")}
}
type grepContentCollector struct {
@@ -494,20 +505,15 @@ func (collector *grepContentCollector) result() Result {
if truncated {
output += fmt.Sprintf("\n\n[truncated: showing first %d matches; narrow path/glob/pattern or increase head_limit]", len(lines))
}
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "narrow path/glob/pattern or increase head_limit")
meta := outputBudgetMeta(budgeted)
if truncated || budgeted.Truncated {
meta := map[string]string{}
if truncated {
meta["truncated"] = "true"
if budgeted.Truncated {
meta["truncation_reason"] = "byte_budget"
} else {
meta["truncation_reason"] = "head_limit"
}
meta["truncation_reason"] = "head_limit"
}
return Result{
Status: StatusOK,
Output: budgeted.Output,
Truncated: truncated || budgeted.Truncated,
Output: output,
Truncated: truncated,
Meta: meta,
}
}
+16 -6
View File
@@ -15,6 +15,10 @@ type listDirectoryTool struct {
scope PathScope
}
func (listDirectoryTool) outputCategory(map[string]any) outputCategory {
return outputCategorySearch
}
func NewListDirectoryTool(workspaceRoot string) Tool {
return NewScopedListDirectoryTool(workspaceRoot, nil)
}
@@ -42,6 +46,14 @@ func NewScopedListDirectoryTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool listDirectoryTool) Run(_ context.Context, args map[string]any) Result {
return tool.run(args, true)
}
func (tool listDirectoryTool) RunWithOptions(_ context.Context, args map[string]any, _ RunOptions) Result {
return tool.run(args, false)
}
func (tool listDirectoryTool) run(args map[string]any, directBudget bool) Result {
// Optional with a "." default: treat an explicit empty path (a common
// weak-model quirk) the same as the key being absent rather than erroring.
requestedPath, err := aliasedStringArg(args, []string{"path", "directory", "dir"}, ".", false, true)
@@ -76,13 +88,11 @@ func (tool listDirectoryTool) Run(_ context.Context, args map[string]any) Result
return okResult("Directory is empty: " + relativePath)
}
output := "Contents of " + relativePath + ":\n\n" + strings.Join(entries, "\n")
budgeted := applyOutputBudget(output, searchOutputBudgetBytes, "use path, recursive=false, or a smaller max_depth to narrow the listing")
meta := outputBudgetMeta(budgeted)
if budgeted.Truncated {
meta["truncated"] = "true"
meta["truncation_reason"] = "byte_budget"
result := Result{Status: StatusOK, Output: output}
if directBudget {
return applyLegacyByteBudgetToResult(result, searchOutputBudgetBytes, "use path, recursive=false, or a smaller max_depth to narrow the listing")
}
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
return result
}
func listDirectoryEntries(path string, depth int, maxDepth int) ([]string, error) {
+254
View File
@@ -0,0 +1,254 @@
package tools
import (
"strconv"
"strings"
)
const (
outputBudgetCategoryMeta = "output_budget_category"
outputBudgetOriginalBytesMeta = "output_budget_original_bytes"
outputBudgetRetainedBytesMeta = "output_budget_retained_bytes"
outputBudgetEstimatedOriginalTokensMeta = "output_budget_estimated_original_tokens"
outputBudgetEstimatedRetainedTokensMeta = "output_budget_estimated_retained_tokens"
outputBudgetReasonMeta = "output_budget_reason"
outputBudgetSpillCreatedMeta = "output_budget_spill_created"
)
// applyRegistryOutputBudget is the common post-redaction semantic budgeting
// boundary for tools that do not already own a deliberate output budget.
func applyRegistryOutputBudget(tool Tool, toolName string, args map[string]any, result Result) Result {
budget := registryOutputBudget(toolName)
if budget.maxEstimatedTokens <= 0 && budget.hardMaxBytes <= 0 {
return result // preserve ZERO_TOOL_OUTPUT_CEILING_TOKENS=0 semantics
}
category := resolveOutputCategory(tool, toolName, args)
budgeted := budgetSemanticOutput(result.Output, category, budget)
if budgeted.truncated {
budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
}
if result.Truncated {
budgeted.truncated = true
if budgeted.reason == "" {
budgeted.reason = result.Meta["truncation_reason"]
if budgeted.reason == "" {
budgeted.reason = "upstream_tool_budget"
}
}
if budgeted.spillPath == "" {
budgeted.spillPath = result.Meta["spill_path"]
}
}
result.Output = budgeted.text
result.Truncated = result.Truncated || budgeted.truncated
result.Meta = addOutputBudgetMetadata(result.Meta, budgeted)
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
// newly combined result so hooks cannot bypass the established safety ceiling.
func (registry *Registry) RebudgetAfterHook(toolName string, args map[string]any, result Result) Result {
result = scrubResultSecrets(result)
tool, _ := registry.Get(toolName)
if _, ok := tool.(selfBudgeting); ok {
// Match the primary registry boundary: self-managed tools keep their
// call-specific capture/output budget instead of being tightened or
// loosened to the generic registry ceiling after hook feedback.
return applySelfManagedOutputBudget(tool, toolName, args, result)
}
result = applyRegistryOutputBudget(tool, toolName, args, result)
return enforceOutputCeiling(toolName, result)
}
func registryOutputBudget(toolName string) outputBudget {
switch toolName {
case "read_file", "read_minified_file":
return outputBudget{maxEstimatedTokens: readOutputBudgetBytes / 4, hardMaxBytes: readOutputBudgetBytes}
case "grep", "glob", "list_directory":
return outputBudget{maxEstimatedTokens: searchOutputBudgetBytes / 4, hardMaxBytes: searchOutputBudgetBytes}
default:
ceilingTokens := resolveOutputCeilingTokens()
if ceilingTokens <= 0 {
return outputBudget{}
}
return outputBudget{maxEstimatedTokens: ceilingTokens, hardMaxBytes: ceilingTokens * 4}
}
}
func resolveOutputCategory(tool Tool, toolName string, args map[string]any) outputCategory {
if provider, ok := tool.(outputPolicyProvider); ok {
if category := provider.outputCategory(args); category != "" {
return category
}
}
switch toolName {
case "Task", "swarm_collect":
return outputCategoryWorker
case "apply_patch":
return outputCategoryDiff
case "write_stdin":
return outputCategoryProcess
default:
return outputCategoryDefault
}
}
// annotateSelfBudgetedOutput records the same compact metadata for tools whose
// existing capture/budget implementation remains authoritative in PR11. It
// does not re-budget their text or claim that raw_bytes represents every byte a
// subprocess produced beyond its established capture limits.
func annotateSelfBudgetedOutput(tool Tool, toolName string, args map[string]any, result Result) Result {
retainedBytes := len(result.Output)
originalBytes := retainedBytes
if parsed, err := strconv.Atoi(result.Meta["raw_bytes"]); err == nil && parsed > originalBytes {
originalBytes = parsed
}
retainedTokens := estimateOutputTokens(result.Output)
originalTokens := retainedTokens
if originalBytes > retainedBytes {
originalTokens = max(originalTokens, estimatedTokensFromBytes(originalBytes))
}
reason := result.Meta["truncation_reason"]
if result.Truncated && reason == "" {
reason = "upstream_tool_budget"
}
observed := budgetedOutput{
text: result.Output,
originalBytes: originalBytes,
retainedBytes: retainedBytes,
estimatedOriginalTokens: originalTokens,
estimatedRetainedTokens: retainedTokens,
truncated: result.Truncated,
category: resolveOutputCategory(tool, toolName, args),
reason: reason,
spillPath: result.Meta["spill_path"],
}
result.Meta = addOutputBudgetMetadata(result.Meta, observed)
return result
}
func shellOutputCategory(command string) outputCategory {
normalized := strings.ToLower(strings.TrimSpace(command))
if containsAny(normalized,
"go test", "pytest", "python -m pytest", "cargo test", "npm test", "npm run test",
"pnpm test", "yarn test", "bun test", "dotnet test", "mvn test", "gradle test", "phpunit") {
return outputCategoryTest
}
if containsAny(normalized, "git diff", "git show", "git format-patch", " diff -u", "diff --git") || strings.HasPrefix(normalized, "diff ") {
return outputCategoryDiff
}
return outputCategoryProcess
}
// attachExistingSpill reuses Zero's hardened per-user spill directory. output
// is the already-redacted text received by this layer; it may itself be a
// capture-bounded view produced by a subprocess tool.
func attachExistingSpill(toolName, output string, budget outputBudget, current budgetedOutput) budgetedOutput {
path := spillTruncatedOutput(toolName, output)
if path == "" {
return current
}
notice := "[zero] full output received by budgeting layer saved to " + path + " (grep or read_file it instead of re-running)"
reduced := outputBudget{
maxEstimatedTokens: max(1, budget.maxEstimatedTokens-estimateOutputTokens("\n"+notice)),
hardMaxBytes: max(1, budget.hardMaxBytes-len("\n"+notice)),
}
base := budgetSemanticOutput(output, current.category, reduced)
text := strings.TrimRight(base.text, "\n") + "\n" + notice
if !fitsOutputBudget(text, budget) {
// An unusually long temp path or tiny configured ceiling can leave no
// room for the full notice. Keep the safe bounded result; the spill still
// exists but is intentionally not advertised with a chopped reference.
current.spillPath = path
return current
}
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
}
func addOutputBudgetMetadata(meta map[string]string, output budgetedOutput) map[string]string {
if meta == nil {
meta = map[string]string{}
}
meta[outputBudgetCategoryMeta] = string(output.category)
meta[outputBudgetOriginalBytesMeta] = strconv.Itoa(output.originalBytes)
meta[outputBudgetRetainedBytesMeta] = strconv.Itoa(output.retainedBytes)
meta[outputBudgetEstimatedOriginalTokensMeta] = strconv.Itoa(output.estimatedOriginalTokens)
meta[outputBudgetEstimatedRetainedTokensMeta] = strconv.Itoa(output.estimatedRetainedTokens)
meta[outputBudgetSpillCreatedMeta] = strconv.FormatBool(output.spillPath != "")
// These established fields describe the final model-visible output. Refresh
// them on every pass because redaction and after-tool hooks may change the
// text even when this pass does not newly truncate it.
meta["emitted_bytes"] = strconv.Itoa(output.retainedBytes)
meta["estimated_tokens"] = strconv.Itoa(output.estimatedRetainedTokens)
if output.reason != "" {
meta[outputBudgetReasonMeta] = output.reason
}
if output.truncated {
// Preserve the existing metadata vocabulary used by callers and tests.
if _, exists := meta["raw_bytes"]; !exists {
meta["raw_bytes"] = strconv.Itoa(output.originalBytes)
}
meta["truncated"] = "true"
meta["truncation_reason"] = output.reason
}
if output.spillPath != "" {
meta["spill_path"] = output.spillPath
}
return meta
}
+381
View File
@@ -0,0 +1,381 @@
package tools
import (
"context"
"fmt"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"unicode/utf8"
)
func TestRegistryBudgetRunsAfterRedactionAndSpillsRedactedOutput(t *testing.T) {
setTestTempDir(t)
secret := "ghp_" + strings.Repeat("a", 36)
big := "HEAD\n" + secret + "\n" + strings.Repeat("🙂 noisy output\n", 20_000) + "TAIL"
registry := NewRegistry()
registry.Register(newCeilingFakeTool("redacted_big", big))
result := registry.Run(context.Background(), "redacted_big", map[string]any{})
if !result.Truncated || result.Meta[outputBudgetSpillCreatedMeta] != "true" {
t.Fatalf("unexpected budget result: truncated=%t meta=%#v", result.Truncated, result.Meta)
}
if strings.Contains(result.Output, secret) || !utf8.ValidString(result.Output) {
t.Fatalf("exposed output leaked secret or invalid UTF-8: %q", result.Output)
}
spillPath := result.Meta["spill_path"]
content, err := os.ReadFile(spillPath)
if err != nil {
t.Fatalf("read spill: %v", err)
}
if strings.Contains(string(content), secret) {
t.Fatal("spill contains unredacted secret")
}
if !strings.Contains(string(content), "HEAD") || !strings.Contains(string(content), "TAIL") {
t.Fatal("spill does not contain the complete redacted output received by the budget layer")
}
}
func TestRegistryBudgetSpillFailureFallsBackToBoundedOutput(t *testing.T) {
temp := t.TempDir()
blockedTemp := filepath.Join(temp, "not-a-directory")
if err := os.WriteFile(blockedTemp, []byte("x"), 0o600); err != nil {
t.Fatal(err)
}
t.Setenv("TMPDIR", blockedTemp)
t.Setenv("TMP", blockedTemp)
t.Setenv("TEMP", blockedTemp)
t.Setenv(outputCeilingEnv, "100")
registry := NewRegistry()
registry.Register(newCeilingFakeTool("spill_failure", strings.Repeat("large output\n", 1000)))
result := registry.Run(context.Background(), "spill_failure", map[string]any{})
if !result.Truncated || len(result.Output) > 400 {
t.Fatalf("fallback is not bounded: truncated=%t bytes=%d", result.Truncated, len(result.Output))
}
if result.Meta[outputBudgetSpillCreatedMeta] != "false" || result.Meta["spill_path"] != "" {
t.Fatalf("spill failure incorrectly advertised: %#v", result.Meta)
}
}
func TestRegistryMissingPolicyUsesDefaultAndExactHardCeiling(t *testing.T) {
setTestTempDir(t)
t.Setenv(outputCeilingEnv, "80")
input := "HEAD\n" + strings.Repeat("x", 3000) + "\nTAIL"
registry := NewRegistry()
registry.Register(newCeilingFakeTool("no_policy", input))
result := registry.Run(context.Background(), "no_policy", map[string]any{})
if got := result.Meta[outputBudgetCategoryMeta]; got != string(outputCategoryDefault) {
t.Fatalf("category = %q, want default", got)
}
if len(result.Output) > 80*4 {
t.Fatalf("output = %d bytes, hard ceiling %d", len(result.Output), 80*4)
}
if result.Meta[outputBudgetOriginalBytesMeta] != strconv.Itoa(len(input)) {
t.Fatalf("original size metadata = %q", result.Meta[outputBudgetOriginalBytesMeta])
}
}
func TestRegistrySmallOutputRemainsByteIdentical(t *testing.T) {
registry := NewRegistry()
registry.Register(newCeilingFakeTool("small_identity", "hello\nworld\n"))
result := registry.Run(context.Background(), "small_identity", map[string]any{})
if result.Output != "hello\nworld\n" || result.Truncated {
t.Fatalf("small output changed: %#v", result)
}
}
func TestDirectReadAndSearchToolsKeepLegacyByteBudgets(t *testing.T) {
root := t.TempDir()
if err := os.WriteFile(filepath.Join(root, "large-read.txt"), []byte(strings.Repeat("source line with enough content\n", 7000)), 0o600); err != nil {
t.Fatal(err)
}
if result := NewReadFileTool(root).Run(context.Background(), map[string]any{"path": "large-read.txt"}); !result.Truncated || len(result.Output) > readOutputBudgetBytes {
t.Fatalf("direct read_file output is not bounded: %#v", result)
}
for index := 0; index < 800; index++ {
name := fmt.Sprintf("%04d-%s.txt", index, strings.Repeat("n", 80))
if err := os.WriteFile(filepath.Join(root, name), []byte("match\n"), 0o600); err != nil {
t.Fatal(err)
}
}
if err := os.WriteFile(filepath.Join(root, "large-content.txt"), []byte(strings.Repeat("match "+strings.Repeat("x", 2000)+"\n", 50)), 0o600); err != nil {
t.Fatal(err)
}
for name, result := range map[string]Result{
"list_directory": NewListDirectoryTool(root).Run(context.Background(), map[string]any{"recursive": false}),
"glob": NewGlobTool(root).Run(context.Background(), map[string]any{"pattern": "*.txt", "limit": 1000}),
"grep files": NewGrepTool(root).Run(context.Background(), map[string]any{"pattern": "match", "output_mode": "files_with_matches"}),
"grep content": NewGrepTool(root).Run(context.Background(), map[string]any{"pattern": "match", "path": "large-content.txt", "head_limit": 50}),
} {
if !result.Truncated || len(result.Output) > searchOutputBudgetBytes {
t.Fatalf("direct %s output is not bounded: truncated=%t bytes=%d meta=%#v", name, result.Truncated, len(result.Output), result.Meta)
}
}
}
func TestRegistryLargeFileUsesSemanticFilePolicy(t *testing.T) {
setTestTempDir(t)
root := t.TempDir()
var content strings.Builder
content.WriteString("HEAD_MARK\n")
for line := 0; line < 7000; line++ {
content.WriteString("ordinary source line with enough text to fill the output budget\n")
}
content.WriteString("TAIL_MARK\n")
if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte(content.String()), 0o600); err != nil {
t.Fatal(err)
}
registry := NewRegistry()
registry.Register(NewReadFileTool(root))
result := registry.Run(context.Background(), "read_file", map[string]any{"path": "large.txt"})
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategoryFile) {
t.Fatalf("large file was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
}
if len(result.Output) > readOutputBudgetBytes {
t.Fatalf("large file output = %d bytes, ceiling %d", len(result.Output), readOutputBudgetBytes)
}
for _, want := range []string{"File: large.txt", "HEAD_MARK", "TAIL_MARK"} {
if !strings.Contains(result.Output, want) {
t.Fatalf("large file output missing %q", want)
}
}
}
func TestRegistryLargeSingleLineFileRetainsHeadAndTail(t *testing.T) {
setTestTempDir(t)
root := t.TempDir()
content := "HEAD_SINGLE_LINE_MARK" + strings.Repeat("x", readOutputBudgetBytes*2) + "TAIL_SINGLE_LINE_MARK"
if err := os.WriteFile(filepath.Join(root, "large.min.js"), []byte(content), 0o600); err != nil {
t.Fatal(err)
}
registry := NewRegistry()
registry.Register(NewReadFileTool(root))
result := registry.Run(context.Background(), "read_file", map[string]any{"path": "large.min.js"})
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategoryFile) {
t.Fatalf("single-line file was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
}
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)
}
}
}
func TestRegistryMinifiedFileUpdatesEmittedBytesAfterBudgeting(t *testing.T) {
setTestTempDir(t)
root := t.TempDir()
content := "HEAD_MINIFIED_MARK" + strings.Repeat("x", readOutputBudgetBytes*2) + "TAIL_MINIFIED_MARK"
if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte(content), 0o600); err != nil {
t.Fatal(err)
}
registry := NewRegistry()
registry.Register(NewReadMinifiedFileTool(root))
result := registry.Run(context.Background(), "read_minified_file", map[string]any{"path": "large.txt"})
if !result.Truncated || result.Meta["emitted_bytes"] != strconv.Itoa(len(result.Output)) {
t.Fatalf("minified metadata does not describe final output: truncated=%t emitted=%q bytes=%d", result.Truncated, result.Meta["emitted_bytes"], len(result.Output))
}
for _, want := range []string{"HEAD_MINIFIED_MARK", "TAIL_MINIFIED_MARK"} {
if !strings.Contains(result.Output, want) {
t.Fatalf("minified single-line file output missing %q", want)
}
}
}
func TestRegistryGrepUsesSemanticMultiFileCoverage(t *testing.T) {
setTestTempDir(t)
root := t.TempDir()
for _, name := range []string{"a.txt", "b.txt", "c.txt", "d.txt"} {
body := strings.Repeat("needle "+strings.Repeat(name, 15)+"\n", 350)
if err := os.WriteFile(filepath.Join(root, name), []byte(body), 0o600); err != nil {
t.Fatal(err)
}
}
registry := NewRegistry()
registry.Register(NewGrepTool(root))
result := registry.Run(context.Background(), "grep", map[string]any{
"pattern": "needle",
"path": ".",
"head_limit": 1400,
})
if !result.Truncated || result.Meta[outputBudgetCategoryMeta] != string(outputCategorySearch) {
t.Fatalf("grep was not semantically budgeted: truncated=%t meta=%#v", result.Truncated, result.Meta)
}
for _, name := range []string{"a.txt:", "b.txt:", "c.txt:", "d.txt:"} {
if !strings.Contains(result.Output, name) {
t.Fatalf("search output lost %s coverage", name)
}
}
}
func TestShellOutputCategoryClassification(t *testing.T) {
tests := []struct {
command string
want outputCategory
}{
{"go test ./...", outputCategoryTest},
{"python -m pytest -q", outputCategoryTest},
{"git diff --stat main...HEAD", outputCategoryDiff},
{"git show HEAD", outputCategoryDiff},
{"make build", outputCategoryProcess},
}
for _, test := range tests {
if got := shellOutputCategory(test.command); got != test.want {
t.Errorf("shellOutputCategory(%q) = %q, want %q", test.command, got, test.want)
}
}
}
func TestShellToolOutputCategoryUsesAllCommandAliases(t *testing.T) {
providers := []struct {
name string
provider outputPolicyProvider
}{
{name: "bash", provider: NewBashTool(t.TempDir()).(outputPolicyProvider)},
{name: "exec_command", provider: NewExecCommandTool(t.TempDir(), newExecSessionManager()).(outputPolicyProvider)},
}
for _, provider := range providers {
for _, alias := range []string{"command", "cmd", "script", "shell"} {
t.Run(provider.name+"/"+alias, func(t *testing.T) {
if got := provider.provider.outputCategory(map[string]any{alias: "go test ./..."}); got != outputCategoryTest {
t.Fatalf("category = %q, want test", got)
}
})
}
}
if got := providers[0].provider.outputCategory(map[string]any{"command": "make build", "cmd": "go test ./..."}); got != outputCategoryProcess {
t.Fatalf("bash alias precedence category = %q, want process", got)
}
if got := providers[1].provider.outputCategory(map[string]any{"cmd": "go test ./...", "command": "make build"}); got != outputCategoryTest {
t.Fatalf("exec_command alias precedence category = %q, want test", got)
}
}
func TestRegistryBudgetPreservesExistingTruncationAndRefreshesMetadata(t *testing.T) {
output := "post-hook output"
result := applyRegistryOutputBudget(newCeilingFakeTool("existing_truncation", output), "existing_truncation", map[string]any{}, Result{
Status: StatusOK,
Output: output,
Truncated: true,
Meta: map[string]string{
"emitted_bytes": "999",
"estimated_tokens": "999",
"truncation_reason": "capture_budget",
},
})
if !result.Truncated || result.Meta["truncated"] != "true" || result.Meta[outputBudgetReasonMeta] != "capture_budget" {
t.Fatalf("prior truncation state was lost: %#v", result)
}
if result.Meta["emitted_bytes"] != strconv.Itoa(len(output)) || result.Meta["estimated_tokens"] != strconv.Itoa(estimateOutputTokens(output)) {
t.Fatalf("final output metadata is stale: %#v", result.Meta)
}
if result.Meta[outputBudgetSpillCreatedMeta] != "false" {
t.Fatalf("existing truncation unexpectedly created a new spill: %#v", result.Meta)
}
}
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(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)
}
}
}
func TestRebudgetAfterHookPreservesSelfManagedExecBudget(t *testing.T) {
setTestTempDir(t)
registry := NewRegistry()
registry.Register(NewExecCommandTool(t.TempDir(), newExecSessionManager()))
t.Run("tight call budget remains authoritative", func(t *testing.T) {
args := map[string]any{"cmd": "make build", "max_output_tokens": 10}
result := registry.RebudgetAfterHook(ExecCommandToolName, args, Result{Status: StatusOK, Output: strings.Repeat("x", 1_000)})
if !result.Truncated || len(result.Output) > 40 {
t.Fatalf("post-hook output escaped call budget: truncated=%t bytes=%d meta=%#v", result.Truncated, len(result.Output), result.Meta)
}
})
t.Run("raised call budget is not replaced by generic ceiling", func(t *testing.T) {
args := map[string]any{"cmd": "make build", "max_output_tokens": 20_000}
output := strings.Repeat("x", 70_000)
result := registry.RebudgetAfterHook(ExecCommandToolName, args, Result{Status: StatusOK, Output: output})
if result.Truncated || result.Output != output {
t.Fatalf("post-hook output was incorrectly limited by generic ceiling: truncated=%t bytes=%d", result.Truncated, len(result.Output))
}
})
}
+213 -2
View File
@@ -4,6 +4,7 @@ import (
"fmt"
"strconv"
"strings"
"unicode/utf8"
)
const (
@@ -11,6 +12,167 @@ const (
searchOutputBudgetBytes = 64 * 1024
)
// outputCategory selects the deterministic retention policy used for a tool
// result. It is deliberately internal: categories are an implementation detail
// of the tools boundary, not part of Zero's public tool protocol.
type outputCategory string
const (
outputCategoryDefault outputCategory = "default"
outputCategoryFile outputCategory = "file"
outputCategorySearch outputCategory = "search"
outputCategoryTest outputCategory = "test"
outputCategoryProcess outputCategory = "process"
outputCategoryDiff outputCategory = "diff"
outputCategoryWorker outputCategory = "worker"
)
// outputBudget combines a provider-neutral estimated-token target with the
// existing byte ceiling. The byte ceiling is always authoritative: semantic
// retention may use fewer bytes, but can never use more.
type outputBudget struct {
maxEstimatedTokens int
hardMaxBytes int
}
// budgetedOutput describes the text retained by a semantic output policy. The
// original size is the output received by this layer; it does not claim to be
// every byte an underlying subprocess may have produced before its own capture
// limits were applied.
type budgetedOutput struct {
text string
originalBytes int
retainedBytes int
estimatedOriginalTokens int
estimatedRetainedTokens int
truncated bool
category outputCategory
reason string
spillPath string
}
// outputPolicyProvider optionally assigns a semantic category to a tool call.
// Tools that do not implement it use outputCategoryDefault.
type outputPolicyProvider interface {
outputCategory(args map[string]any) outputCategory
}
// estimateOutputTokens is a deterministic, provider-neutral estimate used only
// for output budgeting. It is not exact provider tokenization. ASCII non-space
// text uses the repository's established four-bytes-per-token approximation;
// each byte of non-ASCII UTF-8 is counted as a token, intentionally
// overestimating multilingual text and emoji rather than letting them bypass a
// budget. Existing hard byte ceilings remain the final safety limit.
func estimateOutputTokens(value string) int {
asciiNonSpace, nonASCIIBytes := outputTokenComponents(value)
return (asciiNonSpace+3)/4 + nonASCIIBytes
}
func outputTokenComponents(value string) (asciiNonSpace int, nonASCIIBytes int) {
for index := 0; index < len(value); {
r, size := utf8.DecodeRuneInString(value[index:])
if r == utf8.RuneError && size == 1 {
// Invalid input is charged conservatively one token per byte. Budget
// slicing itself remains rune-safe for valid UTF-8 tool output.
nonASCIIBytes++
index++
continue
}
if r < utf8.RuneSelf {
switch r {
case ' ', '\t', '\n', '\r', '\f', '\v':
default:
asciiNonSpace++
}
} else {
nonASCIIBytes += size
}
index += size
}
return asciiNonSpace, nonASCIIBytes
}
const defaultSemanticTruncationNotice = "\n[zero] output truncated\n"
// budgetDefaultOutput applies the safe fallback policy: retain a rune-safe head
// and tail around a stable truncation notice. Small output is returned
// byte-identically. Spill creation is intentionally handled by the shared
// boundary integration so every semantic policy uses the existing spill path.
func budgetDefaultOutput(output string, budget outputBudget) budgetedOutput {
result := budgetedOutput{
text: output,
originalBytes: len(output),
retainedBytes: len(output),
estimatedOriginalTokens: estimateOutputTokens(output),
estimatedRetainedTokens: estimateOutputTokens(output),
category: outputCategoryDefault,
}
if fitsOutputBudget(output, budget) {
return result
}
reason := outputBudgetReason(result.originalBytes, result.estimatedOriginalTokens, budget)
maxContentBytes := len(output)
if budget.hardMaxBytes > 0 {
maxContentBytes = min(maxContentBytes, max(0, budget.hardMaxBytes-len(defaultSemanticTruncationNotice)))
}
// Find the largest deterministic head+tail window that satisfies both the
// estimate and the hard byte ceiling. fitsOutputBudget is monotonic as this
// window grows, so binary search avoids repeatedly trimming one rune at a time.
low, high := 0, maxContentBytes
best := ""
if fitsOutputBudget(defaultSemanticTruncationNotice, budget) {
best = defaultSemanticTruncationNotice
}
for low <= high {
window := low + (high-low)/2
candidate := defaultHeadTail(output, window)
if fitsOutputBudget(candidate, budget) {
best = candidate
low = window + 1
} else {
high = window - 1
}
}
result.text = best
result.retainedBytes = len(best)
result.estimatedRetainedTokens = estimateOutputTokens(best)
result.truncated = true
result.reason = reason
return result
}
func defaultHeadTail(output string, contentBytes int) string {
if contentBytes <= 0 {
return defaultSemanticTruncationNotice
}
headBytes := contentBytes / 2
tailBytes := contentBytes - headBytes
return utf8Prefix(output, headBytes) + defaultSemanticTruncationNotice + utf8Suffix(output, tailBytes)
}
func fitsOutputBudget(output string, budget outputBudget) bool {
if budget.hardMaxBytes > 0 && len(output) > budget.hardMaxBytes {
return false
}
return budget.maxEstimatedTokens <= 0 || estimateOutputTokens(output) <= budget.maxEstimatedTokens
}
func outputBudgetReason(originalBytes int, estimatedTokens int, budget outputBudget) string {
overTokens := budget.maxEstimatedTokens > 0 && estimatedTokens > budget.maxEstimatedTokens
overBytes := budget.hardMaxBytes > 0 && originalBytes > budget.hardMaxBytes
switch {
case overTokens && overBytes:
return "token_and_byte_budget"
case overTokens:
return "estimated_token_budget"
default:
return "hard_byte_ceiling"
}
}
type outputBudgetResult struct {
Output string
Truncated bool
@@ -56,6 +218,7 @@ func estimatedTokensFromBytes(bytes int) int {
type outputBudgetBuilder struct {
builder strings.Builder
capture *boundedBuffer
rawBytes int
maxBytes int
hint string
@@ -65,17 +228,60 @@ 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.maxBytes <= 0 || builder.builder.Len() >= builder.maxBytes {
if builder.capture != nil {
_, _ = builder.capture.Write([]byte(value))
return
}
if builder.maxBytes <= 0 {
builder.builder.WriteString(value)
return
}
if builder.builder.Len() >= builder.maxBytes {
return
}
remaining := builder.maxBytes - builder.builder.Len()
builder.builder.WriteString(utf8Prefix(value, remaining))
}
// applyLegacyByteBudgetToResult preserves direct Tool.Run behavior for callers
// that intentionally bypass Registry.RunWithOptions. Agent/MCP execution uses
// the registry's post-redaction semantic boundary instead.
func applyLegacyByteBudgetToResult(result Result, maxBytes int, hint string) Result {
budgeted := applyOutputBudget(result.Output, maxBytes, hint)
result.Output = budgeted.Output
result.Truncated = result.Truncated || budgeted.Truncated
if result.Meta == nil {
result.Meta = map[string]string{}
}
for key, value := range outputBudgetMeta(budgeted) {
result.Meta[key] = value
}
if budgeted.Truncated {
result.Meta["truncated"] = "true"
result.Meta["truncation_reason"] = "byte_budget"
}
return result
}
func (builder *outputBudgetBuilder) Result() outputBudgetResult {
output := builder.builder.String()
if builder.capture != nil {
output = builder.capture.retained()
}
result := outputBudgetResult{
Output: output,
RawBytes: builder.rawBytes,
@@ -90,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
+90
View File
@@ -0,0 +1,90 @@
package tools
import (
"strings"
"testing"
"unicode/utf8"
)
func TestEstimateOutputTokensASCIIAndUnicode(t *testing.T) {
if got := estimateOutputTokens("abcd efgh"); got != 2 {
t.Fatalf("ASCII estimate = %d, want 2", got)
}
unicodeText := "🙂界"
if got := estimateOutputTokens(unicodeText); got < len([]byte(unicodeText)) {
t.Fatalf("Unicode estimate = %d, want conservative estimate >= %d", got, len([]byte(unicodeText)))
}
if first, second := estimateOutputTokens(unicodeText), estimateOutputTokens(unicodeText); first != second {
t.Fatalf("estimator is not deterministic: %d != %d", first, second)
}
if input := strings.Repeat("\u0080", 4); estimateOutputTokens(input) < len([]byte(input)) {
t.Fatalf("U+0080 estimate = %d, want conservative estimate >= %d", estimateOutputTokens(input), len([]byte(input)))
}
}
func TestBudgetDefaultOutputLeavesSmallOutputByteIdentical(t *testing.T) {
input := "small output\nwith spacing\n"
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 100, hardMaxBytes: 1024})
if got.text != input {
t.Fatalf("small output changed: got %q want %q", got.text, input)
}
if got.truncated || got.originalBytes != len(input) || got.retainedBytes != len(input) {
t.Fatalf("unexpected small-output metadata: %#v", got)
}
}
func TestBudgetDefaultOutputKeepsUTF8HeadTailWithinHardCeiling(t *testing.T) {
input := "HEAD🙂\n" + strings.Repeat("界", 200) + "\nTAIL🙂"
const hardMax = 120
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 10_000, hardMaxBytes: hardMax})
if !got.truncated {
t.Fatal("large output was not truncated")
}
if len(got.text) > hardMax {
t.Fatalf("retained %d bytes, hard ceiling %d", len(got.text), hardMax)
}
if !utf8.ValidString(got.text) {
t.Fatalf("budgeted output is invalid UTF-8: %q", got.text)
}
for _, want := range []string{"HEAD", "TAIL", "output truncated"} {
if !strings.Contains(got.text, want) {
t.Fatalf("budgeted output missing %q: %q", want, got.text)
}
}
}
func TestBudgetDefaultOutputHonorsEstimatedTokenBudget(t *testing.T) {
input := strings.Repeat("🙂", 100)
got := budgetDefaultOutput(input, outputBudget{maxEstimatedTokens: 60, hardMaxBytes: 1024})
if !got.truncated || got.reason != "estimated_token_budget" {
t.Fatalf("unexpected token-budget result: %#v", got)
}
if got.estimatedRetainedTokens > 60 {
t.Fatalf("retained estimate = %d, budget 60", got.estimatedRetainedTokens)
}
}
func TestBudgetDefaultOutputDeterministic(t *testing.T) {
input := "head\n" + strings.Repeat("same line\n", 1000) + "tail\n"
budget := outputBudget{maxEstimatedTokens: 80, hardMaxBytes: 512}
want := budgetDefaultOutput(input, budget)
for iteration := 0; iteration < 20; iteration++ {
got := budgetDefaultOutput(input, budget)
if got != want {
t.Fatalf("iteration %d differs:\n got %#v\nwant %#v", iteration, got, want)
}
}
}
func TestBudgetDefaultOutputTinyBudgetsNeverExceedTheirCeilings(t *testing.T) {
input := strings.Repeat("oversized output ", 100)
for _, budget := range []outputBudget{
{hardMaxBytes: len(defaultSemanticTruncationNotice) - 1},
{maxEstimatedTokens: estimateOutputTokens(defaultSemanticTruncationNotice) - 1},
} {
got := budgetDefaultOutput(input, budget)
if !got.truncated || !fitsOutputBudget(got.text, budget) {
t.Fatalf("tiny budget result does not fit: budget=%#v result=%#v", budget, got)
}
}
}
+17 -15
View File
@@ -25,23 +25,19 @@ 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. Each of these applies its own budget
// before returning: bash (bashOutputBudgetBytes per stream + spill),
// exec_command (model-raisable token budget + spill), read tools (128 KiB),
// search tools (64 KiB).
func (bashTool) managesOutputBudget() {}
func (execCommandTool) managesOutputBudget() {}
func (readFileTool) managesOutputBudget() {}
func (readMinifiedFileTool) managesOutputBudget() {}
func (grepTool) managesOutputBudget() {}
func (globTool) managesOutputBudget() {}
func (listDirectoryTool) managesOutputBudget() {}
// 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() {}
func resolveOutputCeilingTokens() int {
raw := strings.TrimSpace(os.Getenv(outputCeilingEnv))
@@ -61,6 +57,12 @@ func resolveOutputCeilingTokens() int {
// and the spill file agree on what was hidden.
func enforceOutputCeiling(toolName string, result Result) Result {
ceiling := resolveOutputCeilingTokens()
switch toolName {
case "read_file", "read_minified_file":
ceiling = readOutputBudgetBytes / 4
case "grep", "glob", "list_directory":
ceiling = searchOutputBudgetBytes / 4
}
if ceiling <= 0 {
return result
}
+7 -7
View File
@@ -55,10 +55,10 @@ func TestOutputCeilingCapsUnbudgetedTool(t *testing.T) {
if result.Meta["raw_bytes"] != strconv.Itoa(len(big)) {
t.Fatalf("raw_bytes = %s, want %d", result.Meta["raw_bytes"], len(big))
}
if !strings.Contains(result.Output, "full output saved to ") {
if !strings.Contains(result.Output, "full output received by budgeting layer saved to ") {
t.Fatalf("ceiling truncation must include a spill hint: %q", result.Output[:200])
}
start := strings.Index(result.Output, "full output saved to ") + len("full output saved to ")
start := strings.Index(result.Output, "full output received by budgeting layer saved to ") + len("full output received by budgeting layer saved to ")
end := strings.Index(result.Output[start:], " (grep")
content, err := os.ReadFile(result.Output[start : start+end])
if err != nil {
@@ -123,11 +123,6 @@ func TestSelfBudgetingExemptionList(t *testing.T) {
exempt := []Tool{
NewBashTool(dir),
NewExecCommandTool(dir, newExecSessionManager()),
NewReadFileTool(dir),
NewReadMinifiedFileTool(dir),
NewGrepTool(dir),
NewGlobTool(dir),
NewListDirectoryTool(dir),
}
for _, tool := range exempt {
if _, ok := tool.(selfBudgeting); !ok {
@@ -137,4 +132,9 @@ func TestSelfBudgetingExemptionList(t *testing.T) {
if _, ok := NewWebFetchTool().(selfBudgeting); ok {
t.Error("web_fetch must NOT be exempt — the ceiling is its only budget")
}
for _, tool := range []Tool{NewReadFileTool(dir), NewReadMinifiedFileTool(dir), NewGrepTool(dir), NewGlobTool(dir), NewListDirectoryTool(dir)} {
if _, ok := tool.(selfBudgeting); ok {
t.Errorf("%s must use the shared post-redaction budget boundary", tool.Name())
}
}
}
+447
View File
@@ -0,0 +1,447 @@
package tools
import (
"fmt"
"sort"
"strconv"
"strings"
)
// budgetSemanticOutput applies the deterministic policy for category. The
// caller supplies output after secret redaction. Small results pass through
// byte-identically; only oversized results are semantically reduced.
func budgetSemanticOutput(output string, category outputCategory, budget outputBudget) budgetedOutput {
if category == "" {
category = outputCategoryDefault
}
if fitsOutputBudget(output, budget) {
return unchangedBudgetedOutput(output, category)
}
var retained string
switch category {
case outputCategoryFile:
retained = budgetFileLines(output, budget)
case outputCategorySearch:
retained = budgetSearchLines(output, budget)
case outputCategoryTest:
retained = budgetTestLines(output, budget)
case outputCategoryProcess:
retained = budgetProcessLines(output, budget)
case outputCategoryDiff:
retained = budgetDiffUnits(output, budget)
case outputCategoryWorker:
retained = budgetWorkerLines(output, budget)
default:
return budgetDefaultOutput(output, budget)
}
if retained == "" || !fitsOutputBudget(retained, budget) {
fallback := budgetDefaultOutput(output, budget)
fallback.category = category
return fallback
}
return truncatedBudgetedOutput(output, retained, category, "semantic_"+string(category)+"_budget")
}
func unchangedBudgetedOutput(output string, category outputCategory) budgetedOutput {
estimated := estimateOutputTokens(output)
return budgetedOutput{
text: output,
originalBytes: len(output),
retainedBytes: len(output),
estimatedOriginalTokens: estimated,
estimatedRetainedTokens: estimated,
category: category,
}
}
func truncatedBudgetedOutput(original, retained string, category outputCategory, reason string) budgetedOutput {
return budgetedOutput{
text: retained,
originalBytes: len(original),
retainedBytes: len(retained),
estimatedOriginalTokens: estimateOutputTokens(original),
estimatedRetainedTokens: estimateOutputTokens(retained),
truncated: true,
category: category,
reason: reason,
}
}
func budgetFileLines(output string, budget outputBudget) string {
lines := outputLines(output)
priorities := make([]int, 0, len(lines))
// Keep the existing file/range header first, then alternate from the start
// and end so both requested-location boundaries survive.
if len(lines) > 0 {
priorities = append(priorities, 0)
}
for left, right := 1, len(lines)-1; left <= right; left, right = left+1, right-1 {
priorities = append(priorities, left)
if right != left {
priorities = append(priorities, right)
}
}
units := lineDiffUnits(lines)
selected := selectPrioritizedUnits(units, priorities, budget)
for index := range selected {
// Index zero is the file/range header. If no non-blank requested
// content line fits as a complete line, let the common caller fall
// back to its UTF-8-safe head/tail policy rather than returning only
// a header and omission marker for a minified single-line file.
if index > 0 && strings.TrimSpace(lines[index]) != "" {
return renderSelectedUnits(units, selected)
}
}
return ""
}
func budgetSearchLines(output string, budget outputBudget) string {
lines := collapseConsecutiveDuplicateLines(outputLines(output))
priorities := make([]int, 0, len(lines))
for index, line := range lines {
lower := strings.ToLower(line)
// A match body naturally contains the search term (often literally
// "match"); treat only non-result lines as summaries so one busy file
// cannot consume the budget before cross-file representatives are chosen.
if searchResultFile(line) == "" && (strings.Contains(lower, "match") || strings.Contains(lower, "truncated") || strings.Contains(lower, "result")) {
priorities = append(priorities, index)
}
}
if len(lines) > 0 {
priorities = append(priorities, 0, len(lines)-1)
}
seenFiles := map[string]bool{}
for index, line := range lines {
if file := searchResultFile(line); file != "" && !seenFiles[file] {
seenFiles[file] = true
priorities = append(priorities, index)
}
}
priorities = append(priorities, sequence(len(lines))...)
return retainPrioritizedLines(lines, priorities, budget)
}
func budgetTestLines(output string, budget outputBudget) string {
lines := collapseConsecutiveDuplicateLines(outputLines(output))
priorities := make([]int, 0, len(lines))
for index, line := range lines {
if isTestFailureLine(line) {
for contextIndex := max(0, index-2); contextIndex <= min(len(lines)-1, index+3); contextIndex++ {
priorities = append(priorities, contextIndex)
}
}
}
// Final summaries and process status usually live at the tail.
for index := max(0, len(lines)-16); index < len(lines); index++ {
priorities = append(priorities, index)
}
for index := 0; index < min(8, len(lines)); index++ {
priorities = append(priorities, index)
}
priorities = append(priorities, sequence(len(lines))...)
return retainPrioritizedLines(lines, priorities, budget)
}
func budgetProcessLines(output string, budget outputBudget) string {
lines := collapseConsecutiveDuplicateLines(outputLines(output))
priorities := make([]int, 0, len(lines))
for index := 0; index < min(10, len(lines)); index++ {
priorities = append(priorities, index)
}
seenDiagnostic := map[string]bool{}
for index, line := range lines {
lower := strings.ToLower(line)
if containsAny(lower, "error", "warning", "warn:", "fatal", "panic", "failed", "denied") && !seenDiagnostic[line] {
seenDiagnostic[line] = true
priorities = append(priorities, index)
}
}
for index := max(0, len(lines)-16); index < len(lines); index++ {
priorities = append(priorities, index)
}
priorities = append(priorities, sequence(len(lines))...)
return retainPrioritizedLines(lines, priorities, budget)
}
func budgetWorkerLines(output string, budget outputBudget) string {
lines := collapseConsecutiveDuplicateLines(outputLines(output))
priorities := make([]int, 0, len(lines))
for index, line := range lines {
lower := strings.ToLower(line)
if containsAny(lower, "status", "error", "failed", "failure", "session_id", "changed", "files", "tools executed", "conclusion", "result") {
priorities = append(priorities, index)
}
}
// A specialist's final conclusion is conventionally at the end.
for index := max(0, len(lines)-20); index < len(lines); index++ {
priorities = append(priorities, index)
}
for index := 0; index < min(6, len(lines)); index++ {
priorities = append(priorities, index)
}
priorities = append(priorities, sequence(len(lines))...)
return retainPrioritizedLines(lines, priorities, budget)
}
type diffUnit struct {
order int
text string
key bool
}
func budgetDiffUnits(output string, budget outputBudget) string {
units := splitDiffUnits(output)
if len(units) == 0 {
return ""
}
priority := make([]int, 0, len(units))
for index, unit := range units {
if unit.key {
priority = append(priority, index)
}
}
priority = append(priority, sequence(len(units))...)
return retainPrioritizedUnits(units, priority, budget)
}
// splitDiffUnits keeps each hunk indivisible. File headers/stat text are key
// units so broad file coverage is selected before secondary hunks.
func splitDiffUnits(output string) []diffUnit {
lines := outputLines(output)
units := make([]diffUnit, 0)
var current []string
currentKey := false
flush := func() {
if len(current) == 0 {
return
}
units = append(units, diffUnit{order: len(units), text: strings.Join(current, "\n"), key: currentKey})
current = nil
currentKey = false
}
for _, line := range lines {
if strings.HasPrefix(line, "diff --git ") {
flush()
currentKey = true
current = append(current, line)
continue
}
if strings.HasPrefix(line, "@@ ") {
flush()
current = append(current, line)
continue
}
if len(current) == 0 {
currentKey = true // diff stat/summary or standalone file headers
}
current = append(current, line)
}
flush()
return units
}
func retainPrioritizedLines(lines []string, priorities []int, budget outputBudget) string {
return retainPrioritizedUnits(lineDiffUnits(lines), priorities, budget)
}
func lineDiffUnits(lines []string) []diffUnit {
units := make([]diffUnit, 0, len(lines))
for index, line := range lines {
units = append(units, diffUnit{order: index, text: line})
}
return units
}
func retainPrioritizedUnits(units []diffUnit, priorities []int, budget outputBudget) string {
return renderSelectedUnits(units, selectPrioritizedUnits(units, priorities, budget))
}
func selectPrioritizedUnits(units []diffUnit, priorities []int, budget outputBudget) map[int]bool {
selected := map[int]bool{}
indexes := make([]int, 0, len(priorities))
cost := retainedUnitCost{}
for _, index := range priorities {
if index < 0 || index >= len(units) || selected[index] {
continue
}
position := sort.SearchInts(indexes, index)
previous, next := -1, len(units)
if position > 0 {
previous = indexes[position-1]
}
if position < len(indexes) {
next = indexes[position]
}
candidate := cost
candidate.add(units[index].text)
if len(indexes) == 0 {
candidate.addOmission(index)
candidate.addOmission(len(units) - index - 1)
} else {
candidate.removeOmission(next - previous - 1)
candidate.addOmission(index - previous - 1)
candidate.addOmission(next - index - 1)
}
if !candidate.fits(budget) {
continue
}
selected[index] = true
indexes = append(indexes, 0)
copy(indexes[position+1:], indexes[position:])
indexes[position] = index
cost = candidate
}
return selected
}
// retainedUnitCost tracks the exact rendered size of selected units without
// repeatedly rebuilding the full candidate text. Newlines between rendered
// parts add bytes but no estimated tokens, so the token components remain
// additive even as priorities select units out of source order.
type retainedUnitCost struct {
textBytes int
asciiNonSpace int
nonASCIIBytes int
parts int
}
func (cost *retainedUnitCost) add(text string) {
ascii, nonASCII := outputTokenComponents(text)
cost.textBytes += len(text)
cost.asciiNonSpace += ascii
cost.nonASCIIBytes += nonASCII
cost.parts++
}
func (cost *retainedUnitCost) remove(text string) {
ascii, nonASCII := outputTokenComponents(text)
cost.textBytes -= len(text)
cost.asciiNonSpace -= ascii
cost.nonASCIIBytes -= nonASCII
cost.parts--
}
func (cost *retainedUnitCost) addOmission(count int) {
if count > 0 {
cost.add(omittedSectionsMarker(count))
}
}
func (cost *retainedUnitCost) removeOmission(count int) {
if count > 0 {
cost.remove(omittedSectionsMarker(count))
}
}
func (cost retainedUnitCost) fits(budget outputBudget) bool {
bytes := cost.textBytes + max(0, cost.parts-1)
tokens := (cost.asciiNonSpace+3)/4 + cost.nonASCIIBytes
return (budget.hardMaxBytes <= 0 || bytes <= budget.hardMaxBytes) &&
(budget.maxEstimatedTokens <= 0 || tokens <= budget.maxEstimatedTokens)
}
func omittedSectionsMarker(count int) string {
return fmt.Sprintf("[zero] ... %d section(s) omitted ...", count)
}
func renderSelectedUnits(units []diffUnit, selected map[int]bool) string {
indexes := make([]int, 0, len(selected))
for index := range selected {
indexes = append(indexes, index)
}
sort.Ints(indexes)
if len(indexes) == 0 {
return ""
}
parts := make([]string, 0, len(indexes)*2)
previous := -1
for _, index := range indexes {
if previous >= 0 && index != previous+1 {
omitted := index - previous - 1
parts = append(parts, omittedSectionsMarker(omitted))
} else if previous < 0 && index > 0 {
parts = append(parts, omittedSectionsMarker(index))
}
parts = append(parts, units[index].text)
previous = index
}
if previous < len(units)-1 {
parts = append(parts, omittedSectionsMarker(len(units)-previous-1))
}
return strings.Join(parts, "\n")
}
func outputLines(output string) []string {
return strings.Split(strings.TrimRight(output, "\r\n"), "\n")
}
func collapseConsecutiveDuplicateLines(lines []string) []string {
if len(lines) < 2 {
return lines
}
result := make([]string, 0, len(lines))
for index := 0; index < len(lines); {
end := index + 1
for end < len(lines) && lines[end] == lines[index] {
end++
}
result = append(result, lines[index])
if count := end - index; count > 1 {
result = append(result, fmt.Sprintf("[zero] previous line repeated %d more time(s)", count-1))
}
index = end
}
return result
}
func searchResultFile(line string) string {
// Search output is path:line: text. Locate the first numeric field from
// the left after an optional Windows drive prefix. Match text may itself
// contain numeric colon-delimited fields, which are not record structure.
searchFrom := 0
if len(line) >= 2 && line[1] == ':' && ((line[0] >= 'A' && line[0] <= 'Z') || (line[0] >= 'a' && line[0] <= 'z')) {
searchFrom = 2
}
for searchFrom < len(line) {
lineStartOffset := strings.IndexByte(line[searchFrom:], ':')
if lineStartOffset < 0 {
return ""
}
lineStart := searchFrom + lineStartOffset
lineEndOffset := strings.IndexByte(line[lineStart+1:], ':')
if lineEndOffset < 0 {
return ""
}
lineEnd := lineStart + 1 + lineEndOffset
if _, err := strconv.Atoi(line[lineStart+1 : lineEnd]); err == nil {
return strings.TrimSpace(line[:lineStart])
}
searchFrom = lineStart + 1
}
return ""
}
func isTestFailureLine(line string) bool {
lower := strings.ToLower(line)
return containsAny(lower, "--- fail:", "failed", "failure", "panic", "assert", "error", "fatal", "expected", "actual:")
}
func containsAny(value string, needles ...string) bool {
for _, needle := range needles {
if strings.Contains(value, needle) {
return true
}
}
return false
}
func sequence(length int) []int {
result := make([]int, length)
for index := range result {
result[index] = index
}
return result
}
+144
View File
@@ -0,0 +1,144 @@
package tools
import (
"fmt"
"strings"
"testing"
"unicode/utf8"
)
func semanticTestBudget() outputBudget {
return outputBudget{maxEstimatedTokens: 140, hardMaxBytes: 900}
}
func TestFileOutputPolicyKeepsCompleteHeadAndTailLines(t *testing.T) {
lines := []string{"File: huge.go (500 lines)", "1 | package main"}
for index := 2; index < 500; index++ {
lines = append(lines, fmt.Sprintf("%d | value := %d", index, index))
}
lines = append(lines, "500 | // END_MARK")
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategoryFile, semanticTestBudget())
if !got.truncated || !strings.Contains(got.text, "File: huge.go") || !strings.Contains(got.text, "END_MARK") {
t.Fatalf("file policy lost boundaries: %#v", got)
}
if !utf8.ValidString(got.text) || strings.Contains(got.text, "value := 49\ufffd") {
t.Fatalf("file policy produced invalid or split text: %q", got.text)
}
}
func TestSearchOutputPolicyPreservesMultiFileCoverage(t *testing.T) {
var lines []string
for _, file := range []string{"a.go", "b.go", "c.go", "d.go"} {
for line := 1; line <= 30; line++ {
lines = append(lines, fmt.Sprintf("%s:%d: match value %d", file, line, line))
}
}
lines = append(lines, "120 matches found")
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategorySearch, semanticTestBudget())
for _, want := range []string{"a.go:", "b.go:", "c.go:", "d.go:", "120 matches"} {
if !strings.Contains(got.text, want) {
t.Fatalf("search policy missing %q:\n%s", want, got.text)
}
}
}
func TestSearchOutputPolicyPreservesWindowsPathCoverage(t *testing.T) {
var lines []string
for _, file := range []string{"a.go", "b.go", "c.go", "d.go"} {
for line := 1; line <= 30; line++ {
lines = append(lines, fmt.Sprintf(`C:\workspace\%s:%d: match value %d`, file, line, line))
}
}
lines = append(lines, "120 matches found")
got := budgetSemanticOutput(strings.Join(lines, "\n"), outputCategorySearch, semanticTestBudget())
for _, want := range []string{`C:\workspace\a.go:`, `C:\workspace\b.go:`, `C:\workspace\c.go:`, `C:\workspace\d.go:`, "120 matches"} {
if !strings.Contains(got.text, want) {
t.Fatalf("search policy missing %q:\n%s", want, got.text)
}
}
}
func TestSearchResultFileIgnoresNumericMatchContent(t *testing.T) {
tests := []struct {
line string
want string
}{
{line: `C:\workspace\a.go:12: error code:123: failed`, want: `C:\workspace\a.go`},
{line: `/workspace/a.go:12: error code:123: failed`, want: `/workspace/a.go`},
{line: `/workspace/name:with-colon.go:12: match`, want: `/workspace/name:with-colon.go`},
}
for _, test := range tests {
if got := searchResultFile(test.line); got != test.want {
t.Errorf("searchResultFile(%q) = %q, want %q", test.line, got, test.want)
}
}
}
func TestProcessOutputPolicyCollapsesRepetitiveLogsAndKeepsDiagnostics(t *testing.T) {
input := "starting server\n" + strings.Repeat("polling...\n", 300) + "WARNING: queue slow\nERROR: request failed\nshutdown complete"
got := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget())
for _, want := range []string{"starting server", "repeated", "WARNING", "ERROR", "shutdown complete"} {
if !strings.Contains(got.text, want) {
t.Fatalf("process policy missing %q:\n%s", want, got.text)
}
}
if strings.Count(got.text, "polling...") > 1 {
t.Fatalf("repetitive line was not collapsed:\n%s", got.text)
}
}
func TestTestOutputPolicyKeepsFailureAndFinalSummary(t *testing.T) {
input := "=== RUN TestSuite\n" + strings.Repeat("ok progress\n", 300) +
"--- FAIL: TestImportant (0.01s)\n thing_test.go:42: expected 7, got 9\nFAIL\nexit status 1"
got := budgetSemanticOutput(input, outputCategoryTest, semanticTestBudget())
for _, want := range []string{"TestImportant", "expected 7", "FAIL", "exit status 1"} {
if !strings.Contains(got.text, want) {
t.Fatalf("test policy missing %q:\n%s", want, got.text)
}
}
}
func TestDiffOutputPolicyKeepsFilesAndCompleteHunks(t *testing.T) {
var diff strings.Builder
for file := 1; file <= 4; file++ {
fmt.Fprintf(&diff, "diff --git a/f%d.go b/f%d.go\n--- a/f%d.go\n+++ b/f%d.go\n", file, file, file, file)
for hunk := 1; hunk <= 8; hunk++ {
fmt.Fprintf(&diff, "@@ -%d,2 +%d,2 @@\n-old%d_%d\n+new%d_%d\n", hunk, hunk, file, hunk, file, hunk)
}
}
got := budgetSemanticOutput(diff.String(), outputCategoryDiff, outputBudget{maxEstimatedTokens: 220, hardMaxBytes: 1400})
for file := 1; file <= 4; file++ {
if !strings.Contains(got.text, fmt.Sprintf("diff --git a/f%d.go", file)) {
t.Fatalf("diff policy lost file %d header:\n%s", file, got.text)
}
}
for _, line := range strings.Split(got.text, "\n") {
if strings.HasPrefix(line, "@@ ") {
// Each retained hunk is an indivisible three-line unit.
if !strings.Contains(got.text, line+"\n-old") {
t.Fatalf("hunk header was sliced from its body: %q", line)
}
}
}
}
func TestWorkerOutputPolicyKeepsStatusErrorsAndConclusion(t *testing.T) {
input := "session_id: child_1\nstatus: error\n" + strings.Repeat("worker progress\n", 200) +
"tools executed: grep, read_file\nchanged files: a.go, b.go\nerror: verification failed\nConclusion: fix the nil check before merging."
got := budgetSemanticOutput(input, outputCategoryWorker, semanticTestBudget())
for _, want := range []string{"session_id", "status: error", "tools executed", "changed files", "verification failed", "Conclusion"} {
if !strings.Contains(got.text, want) {
t.Fatalf("worker policy missing %q:\n%s", want, got.text)
}
}
}
func TestSemanticPoliciesAreDeterministic(t *testing.T) {
input := "start\n" + strings.Repeat("tick\n", 500) + "ERROR: boom\nend"
want := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget())
for iteration := 0; iteration < 20; iteration++ {
if got := budgetSemanticOutput(input, outputCategoryProcess, semanticTestBudget()); got != want {
t.Fatalf("iteration %d differs:\n got %#v\nwant %#v", iteration, got, want)
}
}
}
+35 -7
View File
@@ -18,6 +18,8 @@ type readFileTool struct {
scope PathScope
}
func (readFileTool) outputCategory(map[string]any) outputCategory { return outputCategoryFile }
func NewReadFileTool(workspaceRoot string) Tool {
return NewScopedReadFileTool(workspaceRoot, nil)
}
@@ -50,10 +52,14 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool readFileTool) Run(ctx context.Context, args map[string]any) Result {
return tool.RunWithOptions(ctx, args, RunOptions{})
return tool.run(args, RunOptions{}, true)
}
func (tool readFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
return tool.run(args, options, false)
}
func (tool readFileTool) run(args map[string]any, options RunOptions, directBudget bool) Result {
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filepath", "filename"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
@@ -86,10 +92,14 @@ func (tool readFileTool) RunWithOptions(_ context.Context, args map[string]any,
// not the authoritative content hash.
options.FileTracker.RecordHash(absolutePath, stats.hash, stats.info)
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines)
maxBytes := 0
if directBudget {
maxBytes = readOutputBudgetBytes
}
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
}
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int) Result {
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int, maxBytes int) Result {
if startLine > total {
return okResult(fmt.Sprintf("File: %s\n(start_line %d is past the end of the file, which has %d lines)", relativePath, startLine, total))
}
@@ -114,7 +124,16 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
header = fmt.Sprintf("File: %s (lines %d-%d of %d)", relativePath, startLine, lastLine, total)
}
budgetedOutput := newOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
// The shared registry boundary applies the file policy after redaction. Its
// input still needs a bounded capture: retain the requested range's head and
// tail instead of materializing an arbitrarily large file in memory. Direct
// Tool.Run calls preserve the legacy prefix-only byte budget.
var budgetedOutput *outputBudgetBuilder
if maxBytes > 0 {
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines to continue with a smaller range")
} else {
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
}
budgetedOutput.WriteString(header)
budgetedOutput.WriteString("\n\n")
if err := appendReadFileRange(budgetedOutput, absolutePath, startLine, selectedLines, width); err != nil {
@@ -128,10 +147,19 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
}
budgeted := budgetedOutput.Result()
meta := outputBudgetMeta(budgeted)
if budgeted.Truncated {
meta := map[string]string{}
if maxBytes > 0 || budgeted.Truncated {
for key, value := range outputBudgetMeta(budgeted) {
meta[key] = value
}
}
if truncated || budgeted.Truncated {
meta["truncated"] = "true"
meta["truncation_reason"] = "byte_budget"
if budgeted.Truncated {
meta["truncation_reason"] = "byte_budget"
} else {
meta["truncation_reason"] = "max_lines"
}
}
return Result{
Status: StatusOK,
+15 -9
View File
@@ -16,6 +16,10 @@ type readMinifiedFileTool struct {
scope PathScope
}
func (readMinifiedFileTool) outputCategory(map[string]any) outputCategory {
return outputCategoryFile
}
func NewReadMinifiedFileTool(workspaceRoot string) Tool {
return NewScopedReadMinifiedFileTool(workspaceRoot, nil)
}
@@ -42,10 +46,14 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool readMinifiedFileTool) Run(ctx context.Context, args map[string]any) Result {
return tool.RunWithOptions(ctx, args, RunOptions{})
return tool.run(args, RunOptions{}, true)
}
func (tool readMinifiedFileTool) RunWithOptions(_ context.Context, args map[string]any, options RunOptions) Result {
return tool.run(args, options, false)
}
func (tool readMinifiedFileTool) run(args map[string]any, options RunOptions, directBudget bool) Result {
requestedPath, err := aliasedStringArg(args, []string{"path", "file", "file_path", "filepath", "filename"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for read_minified_file: " + err.Error())
@@ -92,23 +100,21 @@ func (tool readMinifiedFileTool) RunWithOptions(_ context.Context, args map[stri
savedTokens = estimatedTokensFromBytes(savedBytes)
}
output := header + "\n\n" + result.Content
budgeted := applyOutputBudget(output, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
meta := outputBudgetMeta(budgeted)
meta := map[string]string{}
meta["path"] = relativePath
meta["mode"] = result.Language
meta["compacted"] = strconv.FormatBool(result.Applied)
meta["raw_bytes"] = strconv.Itoa(rawBytes)
meta["compact_bytes"] = strconv.Itoa(compactBytes)
meta["emitted_bytes"] = strconv.Itoa(budgeted.EmittedBytes)
meta["emitted_bytes"] = strconv.Itoa(len(output))
meta["raw_lines"] = strconv.Itoa(rawLines)
meta["emitted_lines"] = strconv.Itoa(minLines)
meta["estimated_tokens_saved"] = strconv.Itoa(savedTokens)
if budgeted.Truncated {
meta["truncated"] = "true"
meta["truncation_reason"] = "byte_budget"
toolResult := Result{Status: StatusOK, Output: output, Meta: meta}
if directBudget {
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
}
return Result{Status: StatusOK, Output: budgeted.Output, Truncated: budgeted.Truncated, Meta: meta}
return toolResult
}
// lineCount reports the number of newline-separated lines in s (an empty string
+9 -4
View File
@@ -126,20 +126,25 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
// args/paths) are redacted at the boundary just like tool output. The output
// ceiling runs after the scrub so the transcript and the spill file agree on
// what was hidden.
ceilingExempt := false
selfManagedOutput := false
var tool Tool
var ok bool
defer func() {
result = scrubResultSecrets(result)
if !ceilingExempt {
if selfManagedOutput {
result = applySelfManagedOutputBudget(tool, name, args, result)
} else {
result = applyRegistryOutputBudget(tool, name, args, result)
result = enforceOutputCeiling(name, result)
}
}()
tool, ok := registry.Get(name)
tool, ok = registry.Get(name)
if !ok {
return errorResult(`Error: Unknown tool "` + name + `".`)
}
if _, ok := tool.(selfBudgeting); ok {
ceilingExempt = true
selfManagedOutput = true
}
if rejecter, ok := tool.(PrePermissionRejecter); ok {
if res, rejected := rejecter.RejectBeforePermission(args); rejected {
+28
View File
@@ -117,6 +117,25 @@ func WriteNDJSON(w io.Writer, t *TurnTrace) error {
return err
}
}
// Output budget events stay in tool-result emission order. Sorting them by
// tool/category would destroy correlation with concurrent calls whose results
// are deliberately consumed in original call order.
for _, event := range t.OutputBudgets {
if err := enc.Encode(map[string]any{
"type": "output_budget",
"tool": event.Tool,
"category": event.Category,
"original_bytes": event.OriginalBytes,
"retained_bytes": event.RetainedBytes,
"estimated_original_tokens": event.EstimatedOriginalTokens,
"estimated_retained_tokens": event.EstimatedRetainedTokens,
"truncated": event.Truncated,
"reason": event.Reason,
"spill_created": event.SpillCreated,
}); err != nil {
return err
}
}
return nil
}
@@ -177,6 +196,15 @@ func WriteText(w io.Writer, t *TurnTrace) error {
for _, c := range counters {
write(" %-22s %d\n", c.Name, c.Value)
}
if len(t.OutputBudgets) > 0 {
write("output budgets:\n")
for _, event := range t.OutputBudgets {
write(" tool=%s category=%s bytes=%d/%d tokens=%d/%d truncated=%t reason=%s spill=%t\n",
event.Tool, event.Category, event.RetainedBytes, event.OriginalBytes,
event.EstimatedRetainedTokens, event.EstimatedOriginalTokens,
event.Truncated, event.Reason, event.SpillCreated)
}
}
return firstErr
}
+100
View File
@@ -0,0 +1,100 @@
package trace
import (
"bytes"
"encoding/json"
"io"
"strings"
"testing"
)
func TestOutputBudgetTraceRoundTripContainsNoOutput(t *testing.T) {
recorder := NewRecorder("session", "run", "")
recorder.Start()
event := OutputBudgetEvent{
Tool: "grep",
Category: "search",
OriginalBytes: 10000,
RetainedBytes: 1000,
EstimatedOriginalTokens: 2500,
EstimatedRetainedTokens: 250,
Truncated: true,
Reason: "semantic_search_budget",
SpillCreated: true,
}
recorder.EmitOutputBudget(event)
var encoded bytes.Buffer
if err := WriteNDJSON(&encoded, recorder.Finish()); err != nil {
t.Fatalf("WriteNDJSON: %v", err)
}
const secretOutput = "secret output body"
if strings.Contains(encoded.String(), secretOutput) {
t.Fatal("output_budget trace contains raw secret output")
}
allowedKeys := map[string]bool{
"type": true, "tool": true, "category": true, "original_bytes": true, "retained_bytes": true,
"estimated_original_tokens": true, "estimated_retained_tokens": true, "truncated": true,
"reason": true, "spill_created": true,
}
findUndocumentedKey := func(record map[string]any) string {
for key := range record {
if !allowedKeys[key] {
return key
}
}
return ""
}
decoder := json.NewDecoder(strings.NewReader(encoded.String()))
var outputBudgetRecord map[string]any
for {
var record map[string]any
if err := decoder.Decode(&record); err != nil {
if err == io.EOF {
break
}
t.Fatalf("decode trace record: %v", err)
}
if record["type"] == "output_budget" {
outputBudgetRecord = record
break
}
}
if outputBudgetRecord == nil {
t.Fatal("missing output_budget trace record")
}
if key := findUndocumentedKey(outputBudgetRecord); key != "" {
t.Fatalf("output_budget trace contains undocumented key %q", key)
}
for key := range allowedKeys {
if _, exists := outputBudgetRecord[key]; !exists {
t.Fatalf("output_budget trace missing documented key %q", key)
}
}
recordJSON, err := json.Marshal(outputBudgetRecord)
if err != nil {
t.Fatalf("encode output_budget trace record: %v", err)
}
var serializedEvent OutputBudgetEvent
if err := json.Unmarshal(recordJSON, &serializedEvent); err != nil {
t.Fatalf("decode output_budget trace event: %v", err)
}
if serializedEvent != event {
t.Fatalf("serialized output_budget event = %#v, want %#v", serializedEvent, event)
}
contaminated := make(map[string]any, len(outputBudgetRecord)+1)
for key, value := range outputBudgetRecord {
contaminated[key] = value
}
contaminated["output"] = secretOutput
if key := findUndocumentedKey(contaminated); key == "" {
t.Fatal("trace key validation would allow raw secret output")
}
parsed, err := ReadNDJSON(strings.NewReader(encoded.String()))
if err != nil {
t.Fatalf("ReadNDJSON: %v", err)
}
if len(parsed.OutputBudgets) != 1 || parsed.OutputBudgets[0] != event {
t.Fatalf("unexpected round trip: %#v", parsed.OutputBudgets)
}
}
+21 -1
View File
@@ -116,6 +116,21 @@ func ReadNDJSON(r io.Reader) (*TurnTrace, error) {
SchemaHash: stringField(obj, "schema"),
CompletePrefixHash: stringField(obj, "complete_prefix"),
})
case "output_budget":
if !sawTraceHeader {
return nil, errors.New("parse trace: not a valid NDJSON trace (no type:trace header)")
}
t.OutputBudgets = append(t.OutputBudgets, OutputBudgetEvent{
Tool: stringField(obj, "tool"),
Category: stringField(obj, "category"),
OriginalBytes: int(parseInt64(obj["original_bytes"])),
RetainedBytes: int(parseInt64(obj["retained_bytes"])),
EstimatedOriginalTokens: int(parseInt64(obj["estimated_original_tokens"])),
EstimatedRetainedTokens: int(parseInt64(obj["estimated_retained_tokens"])),
Truncated: boolField(obj, "truncated"),
Reason: stringField(obj, "reason"),
SpillCreated: boolField(obj, "spill_created"),
})
default:
// Unknown event type: tolerate (forward-compat) but only after a
// header has been seen.
@@ -136,12 +151,17 @@ func ReadNDJSON(r io.Reader) (*TurnTrace, error) {
if !sawTraceHeader {
return nil, errors.New("parse trace: non-empty input had no type:trace header")
}
if len(t.Spans) == 0 && len(t.Counters) == 0 && len(t.PrefixHashes) == 0 {
if len(t.Spans) == 0 && len(t.Counters) == 0 && len(t.PrefixHashes) == 0 && len(t.OutputBudgets) == 0 {
return nil, errors.New("parse trace: header present but no spans, counters, or prefix hashes recovered (corrupt or truncated)")
}
return t, nil
}
func boolField(obj map[string]any, key string) bool {
value, _ := obj[key].(bool)
return value
}
// stringField returns obj[key] as a string, or "" if the key is missing or
// the value is not a string. JSON-marshaled trace events always emit
// string fields as JSON strings, so the type assertion is the right
+16
View File
@@ -188,6 +188,21 @@ func (r *Recorder) EmitPrefixHash(p PrefixHash) {
r.tr.PrefixHashes = append(r.tr.PrefixHashes, p)
}
// EmitOutputBudget records one content-free output budgeting decision. Calls
// are retained in transcript emission order, including when the underlying
// tools executed concurrently.
func (r *Recorder) EmitOutputBudget(event OutputBudgetEvent) {
if r == nil {
return
}
r.mu.Lock()
defer r.mu.Unlock()
if r.finished {
return
}
r.tr.OutputBudgets = append(r.tr.OutputBudgets, event)
}
// Finish stamps CompletedAt, derives each span's parent (by interval
// containment) and exclusive time, and returns a snapshot of the trace. Calling
// Finish more than once returns the same snapshot.
@@ -207,6 +222,7 @@ func (r *Recorder) Finish() *TurnTrace {
snap.Spans = append([]Span(nil), r.tr.Spans...)
snap.Counters = append([]Counter(nil), r.tr.Counters...)
snap.PrefixHashes = append([]PrefixHash(nil), r.tr.PrefixHashes...)
snap.OutputBudgets = append([]OutputBudgetEvent(nil), r.tr.OutputBudgets...)
return &snap
}
+27 -11
View File
@@ -79,20 +79,36 @@ type Counter struct {
Value int64 `json:"value"`
}
// OutputBudgetEvent is compact, content-free metadata for one tool result's
// output budgeting decision. It deliberately carries no output text, paths, or
// arguments so tracing cannot become a secret-bearing side channel.
type OutputBudgetEvent struct {
Tool string `json:"tool"`
Category string `json:"category"`
OriginalBytes int `json:"original_bytes"`
RetainedBytes int `json:"retained_bytes"`
EstimatedOriginalTokens int `json:"estimated_original_tokens"`
EstimatedRetainedTokens int `json:"estimated_retained_tokens"`
Truncated bool `json:"truncated"`
Reason string `json:"reason,omitempty"`
SpillCreated bool `json:"spill_created"`
}
// TurnTrace is the finished record for one agent.Run. It is the value
// emitters serialize; it is not mutated after Finish returns a snapshot.
type TurnTrace struct {
SessionID string `json:"session_id"`
RunID string `json:"run_id"`
Profile string `json:"profile,omitempty"`
StartedAt time.Time `json:"started_at"`
FirstVisibleEventAt time.Time `json:"first_visible_event_at,omitempty"`
FirstUsefulActionAt time.Time `json:"first_useful_action_at,omitempty"`
FirstTokenAt time.Time `json:"first_token_at,omitempty"`
CompletedAt time.Time `json:"completed_at"`
Spans []Span `json:"spans"`
Counters []Counter `json:"counters"`
PrefixHashes []PrefixHash `json:"prefix_hashes,omitempty"`
SessionID string `json:"session_id"`
RunID string `json:"run_id"`
Profile string `json:"profile,omitempty"`
StartedAt time.Time `json:"started_at"`
FirstVisibleEventAt time.Time `json:"first_visible_event_at,omitempty"`
FirstUsefulActionAt time.Time `json:"first_useful_action_at,omitempty"`
FirstTokenAt time.Time `json:"first_token_at,omitempty"`
CompletedAt time.Time `json:"completed_at"`
Spans []Span `json:"spans"`
Counters []Counter `json:"counters"`
PrefixHashes []PrefixHash `json:"prefix_hashes,omitempty"`
OutputBudgets []OutputBudgetEvent `json:"output_budgets,omitempty"`
}
// PrefixHash is one fingerprint of the prompt prefix emitted by an agent run.